백준 16398 - 행성 연결

문제

https://www.acmicpc.net/problem/16398

리뷰

$N$개의 행성을 최소의 비용으로 연결하는 MST를 도출하면 되므로 크루스칼로
풀이할 수 있는 문제였다.

문제를 풀이함에 있어 유의할 점은 $N$이 최대 1000이고 가능한 $C$의 최대값이
1억이므로 MST의 비용 합을 int 타입으로 할 경우 오버플로우가 날 수 있다는
점이다.

로직의 시간 복잡도는 크루스칼의 복잡도인 $O(ElogE)$에 수렴하고 간선의 최대
개수는 $N^2-N$(간선 행렬에서 자기 자신에 대한 간선은 설정하지 않았으므로)이므로
$N$이 1000인 최악의 경우에도 1초의 시간 제한을 무난히 통과한다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

import static java.lang.Integer.*;

public class Main {
    static int[] parent;
    static PriorityQueue<Edge> pq = new PriorityQueue<>(Comparator.comparingInt(edge -> edge.w));

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N = parseInt(br.readLine());
        StringTokenizer st;

        for (int i = 1; i <= N; i++) {
            st = new StringTokenizer(br.readLine());

            for (int j = 1; j <= N; j++) {
                int cost = parseInt(st.nextToken());

                if (i == j) continue;

                pq.offer(new Edge(i, j, cost));
            }
        }

        System.out.println(kruskal(N));
        br.close();
    }

    static long kruskal(int N) {
        parent = new int[N + 1];
        Arrays.fill(parent, -1);

        int selectedEdge = 0;
        long cost = 0;

        while (selectedEdge < N - 1) {
            Edge e = pq.poll();
            int r1 = find(e.u);
            int r2 = find(e.v);

            if (r1 == r2) continue;

            if (parent[r1] < parent[r2]) {
                parent[r1] += parent[r2];
                parent[r2] = r1;
            } else {
                parent[r2] += parent[r1];
                parent[r1] = r2;
            }

            cost += e.w;
            selectedEdge++;
        }

        return cost;
    }

    static int find(int u) {
        if (parent[u] < 0) return u;

        return parent[u] = find(parent[u]);
    }

    static class Edge {
        int u, v, w;

        public Edge(int u, int v, int w) {
            this.u = u;
            this.v = v;
            this.w = w;
        }
    }
}

결과