분리집합(Disjoint Set, Union Find)의 구현 & 크루스칼 알고리즘

분리집합 자료구조는 그래프의 형태로 보이는 문제들에서 유용하게 쓰인다.

1) 두 정점을 연결하는 함수 (Union)

2) 두 정점이 연결되어 있는지를 확인하는 함수 (Find)

가 있다.

대표적으로 Minimum Spanning Tree를 구하는 알고리즘 중 하나인 Kruskal Algorithm이 Union-Find를 이용해 구현한다.

어떠한 두 정점을 연결하고자 할 때, 연결되어 있음이 확인될 경우 cycle이 형성된다는 아이디어를 이용한다.

보통 Union-Find는 parent 배열, height 배열을 이용해서 구현한다.

height 배열은 시간복잡도 최적화를 하기 위해 설정한다.

만약 height가 없다면 $union, find$ 모두 $O(n)$인데, height 배열을 이용해 탐색 횟수를 줄여주면 거의 $O(1)$으로 줄어든다.

 

초기화는 다음과 같이 한다.

int parent[n];
int height[n];
for (int i = 0 ; i < n; i++) {
    parent[i] = i;
    height[i] = 0;
}

---

$union(a, b)$를 할 때,

while (parent[a] != a) a = parent[a];
while (parent[b] != b) b = parent[b];

if (a != b) {
    if (height[a] == height[b]) {
        parent[b] = a; height[a]++;
    }
    else if (height[a] > height[b]) parent[b] = a;
    else parent[a] = b;
}

$a$와 $b$가 다를 때 연결을 한다. 만약 $a = b$일 때 연결을 하면 cycle이 형성된다.

---

$find(a, b)$를 할 때,

while (parent[a] != a) a = parent[a];
while (parent[b] != b) b = parent[b];

if (a == b) printf("YES\n");
else printf("NO\n"); 

$a$와 $b$가 같다면 연결이 되어 있다.

---

이 Union-Find 자료구조를 이용해 MST를 구현하면 다음과 같다. (크루스칼 알고리즘)

#include <cstdio>
#include <algorithm>
#include <vector>
using namespace std;

int main () {
    int n, m;
    scanf("%d %d", &n, &m);
    vector<pair<double, pair<int, int>>> edge;

    for (int i = 0; i < m; i++) {
        int a, b; double c;
        scanf("%d %d %lf", &a, &b, &c);
        a--; b--;
        edge.push_back(make_pair(c, make_pair(a, b)));
    }
    sort(edge.begin(), edge.end());

    int parent[n];
    int height[n];
    for (int i = 0 ; i < n; i++) {
        parent[i] = i;
        height[i] = 0;
    }
    double cost = 0;

    for (int i = 0; i < edge.size(); i++) {
        int x = edge[i].second.first;
        int y = edge[i].second.second;
        double w = edge[i].first;
        while (parent[x] != x) x = parent[x];
        while (parent[y] != y) y = parent[y];

        if (x != y) {
            cost += w;
            if (height[x] == height[y]) {
                parent[y] = x;
                height[x]++;
            }
            else if (height[x] < height[y]) parent[x] = y;
            else parent[y] = x;
        }
    }
    printf("%.0lf", cost);
}