[백준] 1717번 - 집합의 표현 (Gold 4)

문제 링크

[백준] 1717번 - 집합의 표현 (Gold 4)

문제 설명

초기에 (n+1)개의 집합 {0}, {1}, $\cdots$ {n}가 주어져 있다.
다음의 두 기능을 수행하는 프로그램을 작성한다.

1. 입력이 0 a b꼴로 들어오면 a가 포함된 집합과 b가 포함된 집합을 합친다.
2. 입력이 1 a b꼴로 들어오면 a와 b가 같은 집합에 포함되어 있는지 여부를 출력한다.

정답 코드 및 설명

유니온 파인드(Union-find) 알고리즘의 예제이다.
이 알고리즘에서는 집합을 트리 구조로 표현한다.
같은 트리에 속하면 같은 집합에 속하는 것이고, 집합의 대표 원소는 루트 노드이다.
집합을 합칠 때는 둘 중 깊이가 얕은 트리의 루트 노드를 깊이가 깊은 쪽의 루트 노드의 자식으로 삼는다.
이를 코드로 구현하면 아래와 같다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class BOJ1717 {
    int n, m;
    BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    Node[] set;

    class Node {
        int n, level;
        Node parent;

        Node(int n) {
            this.n = n;
            this.level = 0;
            setParent(this);
        }

        void setParent(Node parent) {
            this.parent = parent;
        }
    }

    int findSet(int a) { // a가 속한 집합의 대표 원소 찾기
        Node curr = set[a];
        while (curr != curr.parent) {
            curr = curr.parent;
        }
        return curr.n;
    }

    String isInSameSet(int a, int b) {
        int aSet = findSet(a);
        int bSet = findSet(b);
        if (aSet == bSet) return "YES";
        return "NO";
    }

    void union(int a, int b) {
        int aSet = findSet(a);
        int bSet = findSet(b);
        if (aSet == bSet) return;
        int aSetLevel = set[aSet].level;
        int bSetLevel = set[bSet].level;
        if (aSetLevel >= bSetLevel) {
            set[bSet].parent = set[aSet];
            if (aSetLevel == bSetLevel) set[aSet].level++;
            return;
        }
        set[aSet].parent = set[bSet];
    }

    void solution() throws IOException {
        StringTokenizer st = new StringTokenizer(br.readLine());
        StringBuilder sb = new StringBuilder();
        n = Integer.parseInt(st.nextToken());
        m = Integer.parseInt(st.nextToken());
        set = new Node[n + 1];
        for (int i = 0; i <= n; i++) {
            set[i] = new Node(i);
        }
        while (m-- > 0) {
            st = new StringTokenizer(br.readLine());
            int operator = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            if (operator == 0) union(a, b);
            if (operator == 1) sb.append(isInSameSet(a, b)).append('\n');
        }
        System.out.println(sb);
    }

    public static void main(String[] args) throws IOException {
        new BOJ1717().solution();
    }
}

위의 코드에서는 원소마다 객체를 만들어서 비교하지만, 실제로는 더 간단한 방법이 있다.

1. parents[]라는 배열을 만든다.
2. i번째 원소가 루트 노드이면 parents[i]에 -(트리의 깊이)를 넣는다.
3. i번째 원소가 루트 노드가 아니면 parents[i]에 부모 노드를 넣는다.

이와 비슷한 방식으로 유니온 파인드를 구현한 풀이는 백준 20040번 - 사이클 게임 (Gold 4)에 기술하였다.