PS/Java

[백준 1717] 집합의 표현

siyamaki 2022. 7. 19. 13:14

초기에 {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 를 출력해도 된다)


예제 입력

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

예제 출력

NO
NO
YES

Union-Find를 이용하여 해당 집합들이 서로 동일한 그래프인지 확인한다.

최상단 노드가 같으면 두 원소는 같은 집합이다.

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

public class Main {
    public static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    public static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
    public static int[] arr;
    public static int N, M;
    public static void main(String[] args) throws Exception {
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());
        arr = new int[N + 1];
        for(int i = 0; i <= N; i++) {
        // 자기 자신의 인덱스를 가지도록 초기화
            arr[i] = i;
        }
        for(int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());
            int command = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            if(command == 0) {
            // 합집합을 만든다
                union(a, b);
            } else {
            // a와 b가 같은 최상단 노드를 가리키면 YES 아니면 NO
                if(find(a) == find(b)) {
                    bw.write("YES");
                } else {
                    bw.write("NO");
                }
                bw.newLine();
            }
        }
        br.close();
        bw.flush();
        bw.close();
    }

    public static void union(int from, int to) {
        int x = find(from);
        int y = find(to);
        // 찾으려는 두 값의 최상단 노드를 찾는다
        if(x != y) {
        	// 두 최상단 노드가 다르면 연결시킨다
            if(x < y) {
                arr[y] = x;
            } else {
                arr[x] = y;
            }
        }
    }

    public static int find(int target) {
    // 해당 값의 최상단 노드를 재귀적으로 찾는다
        if(target == arr[target]) {
            return target;
        }
        return arr[target] = find(arr[target]);
    }
}