[백준 / JAVA] 1717번, 집합의 표현 (Union-Find)

문제

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

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

---

입력

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

---

출력

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

---

풀이

유니온 파인드에 대해 알고 있다면 빠르게 풀수 있던 문제입니다. (물론 나는 이걸 하면서 알게 되었지만)

 

부모 노드를 만들어 놓고 각 부모가 어떤 노드인지 설정해 주는 작업이라고 생각하면 됩니다.

처음 Parents[i]는 i(본인)을 설정한 상태입니다.

그리고 0을 받으면 두 노드를 union(합집합) 합니다. 저는 작은 node를 부모로 설정하는 작업을 하였습니다.

마지막으로, 1을 받게되면 두 노드가 같은 부모를 갖고 있는지 connected()를 통해 확인하여 Yes와 No를 반환하는 작업을 하게 됩니다.

 

해당 소스 코드입니다.

import java.io.*;
import java.util.*;

/**
 * 문제: 1717_집합의표현
 * 초기에 $n+1$개의 집합 $\\{0\\}, \\{1\\}, \\{2\\}, \\dots , \\{n\\}$이 있다. 여기에 합집합 연산과,
 * 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.
 * 집합을 표현하는 프로그램을 작성하시오.
 * 입력 :
 * 첫째 줄에 $n$, $m$이 주어진다. $m$은 입력으로 주어지는 연산의 개수이다.
 * 다음 $m$개의 줄에는 각각의 연산이 주어진다.
 * 합집합은 $0$ $a$ $b$의 형태로 입력이 주어진다.
 * 이는 $a$가 포함되어 있는 집합과, $b$가 포함되어 있는 집합을 합친다는 의미이다.
 * 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 $1$ $a$ $b$의 형태로 입력이 주어진다.
 * 이는 $a$와 $b$가 같은 집합에 포함되어 있는지를 확인하는 연산이다.
 */
public class Main {
    static int[] parents;

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

        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());

        parents = new int[n+1];
        for (int i = 1; i <= n; i++) {
            parents[i] = i;
        }

        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < m; i++) {
            st = new StringTokenizer(br.readLine());
            int check = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            if (check == 0){
                union(a,b);
            }
            else {
                if (connected(a, b)){
                    sb.append("YES\\n");
                }
                else {
                    sb.append("NO\\n");
                }
            }
        }
        System.out.println(sb);
        br.close();
    }

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

    static void union(int x, int y) {
        int a = find(x);
        int b = find(y);
        if(a > b){
            parents[b] = a;
        }
        else{
            parents[a] = b;
        }
    }

    static boolean connected(int a, int b) {
        return find(a) == find(b);
    }
}