[BOJ/Union-Find] 백준 1717 - 집합의 표현 (Java)

1717 - 집합의 표현

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

문제

초기에 n + 1개의 집합 {0}, {1}, {2}, ..., {n}
합집합 연산과 두 원소가 같은 집합에 포함되어 있는지 확인하는 연산

입력
첫째 줄: n, m(: 입력으로 주어지는 연산의 개수)
m개의 줄: 각각의 연산 
(합집합: 0 a b = a가 포함되어 있는 집합과 b가 포함되어 있는 집합 합친다는 의미
/ 1 a b: 두 원소가 같은 집합에 포함되어 있는지 확인)
출력: 1로 시작하는 입력에 대해서 a와 b가 같은 집합 포함 O -> YES / X -> NO

---

풀이

Union-Find(Disjoint Set, 서로소 집합)

• 여러 개의 원소가 있을 때 이 원소들이 어떤 그룹(집합)에 속하는지 관리하는 알고리즘
• 시간 복잡도: O(1)
  • 경로 압축(Path Compression): 재귀적으로 루트 찾고 모든 경로 루트로 연결
    -> 트리 높이 줄여서 성능 ↑
    • ex. 초기: 1 -> 2 -> 3 -> 4 (루트)
      / find(1) 실행 후 모든 노드가 루트 가리킴: 1 -> 4, 2 -> 4, 3 -> 4, 4 -> 4
  • 랭크 최적화(Union by Rank): 한 집합의 루트를 다른 집합의 루트 밑에 붙임
    • ex. 초기: 1 - 1, 2 - 2, 3 - 3
      / union(1, 2): 1, 2 -> 1, 3 -> 3 / union(2, 3): 1 - 2 - 3
      => 전체가 하나의 그룹이 됨
• `find(v)`: 루트 찾기 / `union(a, b)`: a, b가 속한 집합 하나로 합침 / `unf[]`: 부모 노드 저장
  
  • `unf[i] = i`: 초기에는 모든 원소가 자기 자신이 루트

static int[] unf;

static int find(int v) {
    if (v == unf[v]) return v;
    return unf[v] = find(unf[v]); // 경로 압축
}

static void union(int a, int b) {
    int fa = find(a);
    int fb = find(b);
    if (fa != fb) unf[fa] = fb;
}

---

코드

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

// 집합의 표현
public class boj_1717 {
    static int[] unf;

    static int find(int v) {
        if (v == unf[v]) return v;
        return unf[v] = find(unf[v]);
    }

    static void union(int a, int b) {
        int fa = find(a);
        int fb = find(b);
        if (fa != fb) unf[fa] = fb;
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        StringBuilder sb = new StringBuilder();

        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());

        unf = new int[n + 1];
        for (int i = 1; i <= n; i++) unf[i] = i;

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

            if (op == 0) union(a, b);
            else {
                if (find(a) == find(b)) sb.append("YES").append('\n');
                else sb.append("NO").append('\n');
            }
        }

        System.out.println(sb);
    }
}