백준 1717 - 집합의 표현

문제 링크

Disjoint set의 union, find 연산을 이용하는 문제이다.

Disjoint set을 표현하는데, 실제 구현할 때는 각 원소의 representative가 누군지 저장하는 parent 배열을 이용한다.

초기에는 각 원소가 자기 자신을 representative로 가지고 있게 parent[x] = x 와 같은 형태로 세팅한다.

• \(\{x\}\), \(\{y\}\) 집합을 합치는것은 \(y\)의 representative의 representative을 \(x\)의 representative로 대치하는 것.
• \(x\),\(y\)가 같은 집합에 속해있는지 보는것은 \(x\)의 representative와 \(y\)의 repr가 같은지 보는것.
• repr을 찾는 함수를 find()라 하면, find()의 pseudo 코드는 아래와 같다.

find(x) = 
    if parent[x] = x then return x
    else 
        r = find(parent[x])
        parent[x] = r // path compression
        return r

Path compression

\(\{e_1\}\)을 \(\{e_2\}\)에 \(\{e_2\}\)를 \(\{e_3\}\)에 \(\{e_3\}\)를 \(\{e_4\}\)에 합치는 식으로 계속 가다보면 parent 배열은 다음 그림과 같이 구성된다.

이런 상황에서 \(e_1\)의 representative인 \(e_4\)를 찾으려면 find 함수를 재귀적으로 3번 더 호출해야한다. 매번 이렇게 하면 비효율적이므로, find 함수 호출의 결과값을 가지고 콜스택을 빠져나오면서 원소들의 representative을 다시 업데이트 해준다.

소스코드

#include <stdio.h>

#define N 1000000

int parent[N+1];

int find(int x) {
    if (parent[x] == x) {
        return x;
    }

    int r = find(parent[x]);
    parent[x] = r;
    return r;
}

void merge(int a, int b) {
    int ra = find(a);
    int rb = find(b);

    parent[rb] = ra;
}

int main() {
    int i, n, m;
    for (i=0; i<=N; i++) {
        parent[i] = i;
    }
    scanf("%d %d", &n, &m);
    int x, a, b;
    for (int i=0; i<m; i++) {
        scanf("%d %d %d", &x, &a, &b);
        if (x == 0) {
            merge(a, b);
        } else { 
            printf("%s\n", find(a) == find(b) ? "YES" : "NO");
        }
    }
    return 0;
}