[ 백준 ] 1717 집합의 표현

📌 Problem

초기 집합들은 0 ~ N으로 각각 자신의 번호를 원소로만 한다. 이후 합연산을 진행하여 두 원소가 속한 집합을 합한다. 합연산 외에는 동일 집합인지를 확인하는 연산을 수행하여 동일한 집합에 속해 있으면 YES를 아니면 NO를 출력한다.

📌 Solution

Union-Find을 사용하는 대표적인 문제이다. 그러나 일반적인 방법으로 사용하면 시간 초과 때문에 틀리게 된다. 그래서 find 연산을 할 때마다 root를 업데이트 한다.

📌 Code

public void solve() throws IOException {
    StringTokenizer tokenizer = new StringTokenizer(reader.readLine());
    int[] roots = IntStream.rangeClosed(0, Integer.parseInt(tokenizer.nextToken())).toArray();
    int M = Integer.parseInt(tokenizer.nextToken());
    while (M-- > 0) {
        tokenizer = new StringTokenizer(reader.readLine());
        int type = Integer.parseInt(tokenizer.nextToken());
        int rootA = find(roots, Integer.parseInt(tokenizer.nextToken()));
        int rootB = find(roots, Integer.parseInt(tokenizer.nextToken()));
        if (type == 0) union(roots, rootA, rootB);
        else result.append(rootA == rootB ? "YES" : "NO").append("\n");
    }
}

private int find(int[] roots, int leaf) {
    return roots[leaf] = roots[leaf] == leaf ? leaf : find(roots, roots[leaf]);
}

private void union(int[] roots, int rootA, int rootB) {
    if (rootA <= rootB) roots[rootB] = rootA;
    else roots[rootA] = rootB;
}