[JAVA] 백준 (골드5) 1717번 집합의 표현

링크

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

---

문제 설명

정답률 28.085%
초기에 $n+1$개의 집합 $\{0\}, \{1\}, \{2\}, \dots , \{n\}$이 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

집합을 표현하는 프로그램을 작성하시오.

---

입력 예제

• 첫째 줄에 $n$, $m$이 주어진다.
• $m$은 입력으로 주어지는 연산의 개수이다.
• 다음 $m$개의 줄에는 각각의 연산이 주어진다.
• 합집합은 $0$ $a$ $b$의 형태로 입력이 주어진다.
  • 이는 $a$가 포함되어 있는 집합과, $b$가 포함되어 있는 집합을 합친다는 의미이다.
• 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 $1$ $a$ $b$의 형태로 입력이 주어진다.
  • 이는 $a$와 $b$가 같은 집합에 포함되어 있는지를 확인하는 연산이다.
•  $1 ≤ n ≤ 1,000,000$ 
•  $1 ≤ m ≤ 100,000$ 
•  $0 ≤ a, b ≤ n$ 
•  $a$, $b$는 정수
•  $a$와 $b$는 같을 수도 있다.

7 8
0 1 3
1 1 7
0 7 6
1 7 1
0 3 7
0 4 2
0 1 1
1 1 1

---

출력 예제

• 1로 시작하는 입력에 대해서 $a$와 $b$가 같은 집합에 포함되어 있으면 "YES" 또는 "yes"를, 그렇지 않다면 "NO" 또는 "no"를 한 줄에 하나씩 출력한다.

NO
NO
YES

---

풀이

문제 자체는 어렵지 않은 Union-Find 알고리즘 문제이다. 부모 노드를 찾는 메서드와 노드를 합치는 메서드를 다음과 같이 구현한다.

static int getParents(int x) {
    if (parents[x] == x) {
        return x;
    }
    return getParents(parents[x]);
}

static void unionParents(int a, int b) {
    a = getParents(a);
    b = getParents(b);
    if (a < b) {
        parents[b] = a;
    } else {
        parents[a] = b;
    }
}

하지만 집합의 원소의 최대 개수는 $10^6$이고 연산의 최대 개수도 $10^5$이므로 이렇게 코드를 구현할 경우 시간 초과가 발생한다.

현재 unionParents 메서드로 합집합을 만들고 getParent 메서드로 특정 노드의 부모 노드를 찾을 때 시간 복잡도는 $O(n)$가 된다.

가령 1, 2, 3, 4란 노드가 있을 때 (1, 2), (2, 3), (3, 4) 이렇게 합집합을 구성하면 다음과 같은 트리가 만들어진다.

getParent(4)로 대표 노드인 1을 찾으려면 트리의 높이만큼의 시간 복잡도를 가진다.

이때 경로 압축을 하여 메서드를 수정하면 다음과 같다.

static int getParents(int x) {
    if (parents[x] == x) {
        return x;
    }
    return parents[x] = getParents(parents[x]);
}

수정된 getParents 메서드를 수행하면 연산을 할 때 거치는 노드들이 다음과 같이 대표 노드로 바로 연결된다. 이렇게 되면 추후 노드와 관련된 getParents 메서드의 시간복잡도는 $O(1)$이 된다.

코드

//백준
public class Main {

    static int[] parents;

    public static void main(String[] args) throws IOException {

        System.setIn(new FileInputStream("src/input.txt"));
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int n = Integer.parseInt(st.nextToken());  //10^6
        int m = Integer.parseInt(st.nextToken());  //10^5
        parents = new int[n + 1];
        for (int i = 0; i < n + 1; i++) {
            parents[i] = i;
        }

        StringBuilder sb = new StringBuilder();

        for (int i = 0; i < m; i++) {
            st = new StringTokenizer(br.readLine());
            int x = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());

            if (x == 0) {
                unionParents(a, b);
            } else if (x == 1) {
                if (checkParents(a, b)) {
                    sb.append("YES").append("\n");
                } else {
                    sb.append("NO").append("\n");
                }
            }
        }
        System.out.println(sb);
    }
	//부모 노드를 찾는 메서드
    static int getParents(int x) {
        if (parents[x] == x) {
            return x;
        }
        return parents[x] = getParents(parents[x]);
    }
	//부모 노드를 합치는 메서드
    static void unionParents(int a, int b) {
        a = getParents(a);
        b = getParents(b);

        if (a < b) {
            parents[b] = a;
        } else {
            parents[a] = b;
        }
    }
	//같은 집합 여부를 판단하는 메서드
    static boolean checkParents(int a, int b) {
        a = getParents(a);
        b = getParents(b);

        return a == b;
    }
}