백준 18116 - 로봇 조립

문제

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

리뷰

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

최적화된 유니온 파인드 로직을 사용하여 풀이했으며 두 부품이
주어지는 I 쿼리의 경우 union 연산으로 처리를, Q 쿼리의 경우
find 연산을 통해 루트를 찾아내어 음수로 표현된 자식 노드 수의
값을 구하여 처리했다.

로직의 시간복잡도는 쿼리를 처리하는 부분에서 $O(Na(k))$, $k=1,000,000$의
형태로 수렴하고 $N=10^6$인 최악의 경우에도 무난히 제한 조건 4초를 통과한다.

코드

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 = new int[1_000_001];

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

        Arrays.fill(parent, -1);

        String q;
        int u, v;
        StringBuilder sb = new StringBuilder();
        while (N-- > 0) {
            st = new StringTokenizer(br.readLine());
            q = st.nextToken();

            if (q.equals("I")) {
                u = parseInt(st.nextToken());
                v = parseInt(st.nextToken());

                union(u, v);
            } else {
                u = parseInt(st.nextToken());
                int root = find(u);
                sb.append(-parent[root]).append("\n");
            }
        }

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

    static void union(int u, int v) {
        int r1 = find(u);
        int r2 = find(v);

        if (r1 == r2) return;

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

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

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

결과