백준 1717번 집합의 표현 (JAVA)

백준 1717번 집합의 표현

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

1717번: 집합의 표현

문제

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

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

입력

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

출력

1로 시작하는 입력에 대해서 한 줄에 하나씩 YES/NO로 결과를 출력한다. (yes/no 를 출력해도 된다)

---

소스 코드

- Union-Find (Disjoint Set) 기본 유형

- 부모 집합 구현 및 비교가 핵심

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.StringTokenizer;

public class Main {

    private static int n, m;
    private static int[] parent;
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st;

        st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken()); //n : 집합의 수
        m = Integer.parseInt(st.nextToken()); //m :연산의 수

        parent = new int[n+1];
        //배열 초기화
        for (int i = 1; i <= n; i++) {
            parent[i] = i;
        }

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

            //결과 출력
            if(a == 1){
                boolean equals = equalsparent(b, c);
                if(equals == true) bw.write("YES" + "\n");
                else bw.write("NO" + "\n");
            //a == 0 : 집합 합침
            }else{
                union(b, c);
            }
            bw.flush();
        }
        bw.close();
        br.close();
    }

    //집합 합치기
    private static void union(int a, int b) {
        int x = find(a);
        int y = find(b);

        if (x < y) {
            parent[y] = x;
        } else {
            parent[x] = y;
        }
    }

    //최상위 부모 찾기
    private static int find(int a) {
        if (parent[a] == a)
            return a;
        return parent[a] = find(parent[a]);
    }

    //같은 집합인지 확인
    private static boolean equalsparent(int a, int b) {
        int x = find(a);
        int y = find(b);
        return x == y;
    }
}

Union Find 기본 코드 템플릿 참고

[Union-Find, Disjoint Set] 알고리즘 기본 코드 형식