[백준] 1717번 집합의 표현 - Java

문제 링크

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

문제

초기에
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"를 한 줄에 하나씩 출력한다.

알고리즘: union find, 자료 구조

풀이

union find 알고리즘을 사용하면 쉽게 풀린다.
피연산자 각각 최상위 부모가 누구인지 찾는다. -> find

최상위 부모가 다르다 그러면 같은 집합이 아님을 뜻한다.
최상위 부모가 같으면 같은 집합을 뜻한다.

최상위 부모가 다른데 0연산이라면 두 집합을 합쳐야 한다. 합치는 건 간단하다. 한쪽 최상위 부모 노드에 나머지 한쪽 최상위 부모 노드를 연결하면 된다. -> union

여기서 최상위 부모를 찾는 과정에서 트리가 선형으로 이루어졌을 경우 find 메소드 자체가 매우 무거워진다.
그렇기 때문에 높이 압축을 해줘야 한다.
ex) 5->4->3->2->1 find(5)를 했을 때 선형으로 이루어졌기 때문에 5번의 탐색이 필요
높이 압축을 하면 5->1, 4->1, 3->1, 2->1 이런 식으로 트리가 구성된다.
즉 높이 압축을 한 번만 해놓으면 2번 이상의 탐색이 필요했던 5,4,3 노드가 1번의 탐색만으로 최상위 부모 노드를 찾을 수 있게 된다.

소스 코드

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

public class Main {
    static int N,M;
    static int set[];
    static StringBuilder sb = new StringBuilder();
    public static void main(String args[]) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());
        set = new int[N+1];
        for(int i=0; i<=N; i++) set[i] = i;
        for(int i=0; i<M; i++) {
            StringTokenizer n_st = new StringTokenizer(br.readLine());
            int operation = Integer.parseInt(n_st.nextToken());
            int operand1 = Integer.parseInt(n_st.nextToken());
            int operand2 = Integer.parseInt(n_st.nextToken());
            if(operation == 0) {
                //union
                union(operand1, operand2);
            } else if(operation == 1) {
                //tlp_find
                int tlp1 = tlp_find(operand1);
                int tlp2 = tlp_find(operand2);
                if(tlp1 == tlp2) sb.append("YES\n");
                else sb.append("NO\n");
            }
        }
        System.out.println(sb.toString().trim());
    }
    
    static void union(int oprand1, int oprand2) {
        int tlp_op1 = tlp_find(oprand1);
        int tlp_op2 = tlp_find(oprand2);
        if(tlp_op1 != tlp_op2) {
            if(tlp_op1 < tlp_op2) {
                set[tlp_op2] = tlp_op1;
            } else {
                set[tlp_op1] = tlp_op2;
            }
        }
    }
    
    static int tlp_find(int oprand) {
        if(oprand == set[oprand]) {
            return oprand;
        }
        int v = tlp_find(set[oprand]);
        set[oprand] = v;// 높이 압축
        return v;
    }
}