[BOJ 1717] 집합의 표현 (Java)

BOJ 1717 집합의 표현

난이도 : 골드 4

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

1717번: 집합의 표현

 

전략

크루스칼 알고리즘에서 구현했던 것처럼 parent 배열로 union find를 시도했지만 시간초과가 발생했다.

집합들이 A <- B <- C <- .. 이런 식으로 극단적으로 이어져있는 경우에 root를 만날 때까지 모두 순회하면서 발생하게 된다.

간단한 수정을 통해서 효율을 높일 수 있다.

find 과정에서 parent[i]가 root를 가리키도록 하면 순회의 중복을 줄일 수 있다.

 

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;

public class Main {

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N, M;
        String[] line = br.readLine().split(" ");
        N = Integer.parseInt(line[0]);
        M = Integer.parseInt(line[1]);

        int[] parent = new int[N+1];
        // init
        for (int i=0; i<=N; i++)
            parent[i] = i;

        for (int i=0; i<M; i++) {
            line = br.readLine().split(" ");
            int op = Integer.parseInt(line[0]);
            int set1 = Integer.parseInt(line[1]);
            int set2 = Integer.parseInt(line[2]);

            if (op == 0) {
                union(parent, set1, set2);
            } else if (op == 1) {
                int find1 = find(parent, set1);
                int find2 = find(parent, set2);
                if (find1 == find2)
                    System.out.println("YES");
                else
                    System.out.println("NO");
            }
        }
    }

    static int find(final int[] parent, int i) {
        if (parent[i] == i)
            return i;
        return parent[i] = find(parent, parent[i]);
    }

    static void union(final int[] parent, int a, int b) {
        int parentA = find(parent, a);
        int parentB = find(parent, b);

        if (parentA < parentB)
            parent[parentB] = parentA;
        else
            parent[parentA] = parentB;
    }

}