유니온 파인드 (Union Find) 구현하기

이번 포스트에서는 시간 복잡도 O(α(N))의 유니온 파인드(서로소 집합, Disjoint Set Union) 알고리즘을 구현해보겠습니다.

• 서로소 집합(Disjoint Set): 서로 공통된 원소가 없는 두 개 이상의 집합을 의미합니다.

find()로 루트 노드를 찾고, uni()로 두 집합을 합칩니다.
경로 압축과 트리 높이 기준 병합 덕분에, 거의 상수 시간에 가까운 성능을 보장합니다.

* α(N): 아커만 함수의 역함수로 사실상 상수 시간에 근접

1. p[x]: 상위 노드 저장하는 부모 배열

• p[x] < 0: 루트 노드, 절댓값은 집합의 크기
• p[x] > 0: x의 부모 노드 번호

2. find(x): 루트 노드 찾기 (★ 경로 압축)

int find(int x) {
    if (p[x] < 0) return x;
    return p[x] = find(p[x]); // 경로 압축
}

• x의 루트를 찾을 때까지 재귀 호출
• p[x] = find(p[x])로 루트를 직접 가리키게 업데이트 (★ 최적화1: 경로 압축)

3. uni(u, v): 두 집합 합치기 (★ 유니온 바이 랭크)

bool uni(int u, int v) {
    u = find(u), v = find(v);
    if (u == v) return false;        // 이미 같은 집합
    if (p[u] < p[v]) swap(u, v);     // 항상 u가 부모
    if (p[u] == p[v]) p[u]--;        // 트리 높이 반영
    p[v] = u;
    return true;
}

• union은 C++ 예약어이므로 함수명을 uni로 사용
• 이미 같은 집합이면 바로 false 반환
• 항상 더 작은 트리를 큰 트리에 합침
• 두 루트의 높이가 같을 때만 p[u]--로 높이 갱신

루트의 높이가 같으면 한쪽 트리 전체가 다른 트리 아래로 들어가므로, 루트 노드 뒤에 레벨이 하나 추가되어 높이가 1 증가합니다.

---

대표 문제 풀어보기

[백준 Gold V] 소셜 네트워킹 어플리케이션 - 7511에서 유니온-파인드를 이용하여 두 원소가 같은 집합에 속하는지 확인하는 문제입니다. 또한, 삼성 SW 역량테스트와 유사하게 여러 테스트케이스를 처리하도록 되어 있어 배열 초기화에 유의해야 합니다.

#include <bits/stdc++.h>
using namespace std;

int T, N, K, O;
vector<int> p(1000001, -1);

int find(int x);
bool uni(int u, int v);

int main() {
    cin >> T;
    for (int i = 1; i <= T; i++) {
        cin >> N;
        fill(p.begin(), p.end(), -1);
        cin >> K;
        while(K--) { // 두 원소에 대해 집합 합치기
            int u, v; cin >> u >> v;
            uni(u, v);
        }
        cin >> O;
        cout << "Scenario " << i << ":\n";
        while(O--) { // 두 원소가 같은 집합에 속하는지 출력
            int u, v; cin >> u >> v;
            if (find(u) == find(v)) cout << "1\n";
            else cout << "0\n";
        }
        cout << "\n";
    }
}

int find(int x) {
    if (p[x] < 0) return x;
    return p[x] = find(p[x]);
}

bool uni(int u, int v) {
    u = find(u), v = find(v);
    if (u == v) return false;
    if (p[v] < p[u]) swap(v, u); // v가 크면 swap, 항상 u가 더 우선순위가 높게
    if (p[v] == p[u]) p[u]--;
    p[v] = u;
    return true;
}

초기화 관련 주의사항

여러 테스트케이스를 처리할 때마다 배열을 초기화해 주어야 합니다. 이를 빼먹으면 이전 테스트케이스에서 변경된 값이 다음 테스트케이스에 영향을 미칠 수 있습니다.

fill(p.begin(), p.end(), -1); // 배열 초기화 함수

배열 초기화에 대한 자세한 방법은 다차원 배열 초기화 방법 총정리 (2차원 배열 초기화), 문제를 풀면서 발생한 실수는 vector::resize()로 초기화를 기대하면 안 되는 이유 포스트에 별도로 정리해두었습니다.