백준 17250 - 은하철도

문제

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

리뷰

유니온 파인드를 이용하여 풀이할 수 있는 간단한 문제였다.

유니온 파인드시 사용되는 parent 배열을 $2 \times (N+1)$의 형태로
선언하고, parent[0]에 해당 분리집합에 속하는 은하의 행성 수를
누적하도록 로직을 정의하였다.

유의할 부분은 유니온 연산이 수행되지 않을때(새로 연결된 두 은하가 이미
같은 분리 집합에 속해있을 때)에도 출력 처리를 해주어야 한다는 점이었다.

로직의 시간복잡도는 유니온 파인드를 수행하는 부분의 $O(Ma(N))$으로 수렴하며
$N=M=100,000$인 최악의 경우에도 무난히 제한 조건 5초를 통과한다.

코드

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

import static java.lang.Integer.parseInt;

public class Main {
    static int[][] parent;

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

        int N = parseInt(st.nextToken());
        int M = parseInt(st.nextToken());

        parent = new int[2][N + 1];
        Arrays.fill(parent[1], -1);

        int planetNum;
        for (int i = 1; i <= N; i++) {
            planetNum = parseInt(br.readLine());
            parent[0][i] = planetNum;
        }

        int u, v;
        StringBuilder sb = new StringBuilder();

        while (M-- > 0) {
            st = new StringTokenizer(br.readLine());
            u = parseInt(st.nextToken());
            v = parseInt(st.nextToken());

            int r1 = find(u);
            int r2 = find(v);

            if (r1 == r2) {
                sb.append(parent[0][r1]).append("\n");
                continue;
            }

            int root;

            if (parent[1][r1] < parent[1][r2]) {
                parent[1][r1] += parent[1][r2];
                parent[1][r2] = r1;
                parent[0][r1] = parent[0][r1] + parent[0][r2];

                root = r1;
            } else {
                parent[1][r2] += parent[1][r1];
                parent[1][r1] = r2;
                parent[0][r2] = parent[0][r2] + parent[0][r1];

                root = r2;
            }

            sb.append(parent[0][root]).append("\n");
        }

        System.out.print(sb);
        br.close();
    }


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

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

결과