(C++) Union-find 구현

개요

union-find 구현이다.

• 1-based index
• struct-based design

설명

vector<int> p: 각 원소의 부모 원소를 저장하는 배열로, 초깃값은 자신(p[i] = i)이다.

vector<int> size: 각 원소가 부모인 집합의 크기를 저장하는 배열로, 초깃값은 1이다.

unite 함수: small to large, 작은 집합을 큰 집합으로 합치는 테크닉을 활용해 O(logN)의 시간복잡도를 보장한다.

get 함수: small to large으로 합친 독립 집합에 path compression을 적용해 시간복잡도 amortized O(α(N))을 보장한다.

구현

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

struct DSU {
    vector<int> p, size;
    DSU(int N) : p(N + 1), size(N + 1, 1) {
        iota(p.begin(), p.end(), 0);
    };

    int get(int a) {
        if (p[a] == a) return a;
        return p[a] = get(p[a]);
    }

    bool unite(int a, int b) {
        a = get(a);
        b = get(b);

        if (a == b) return false;

        if (size[a] < size[b]) swap(a, b);
        size[a] += size[b];
        p[b] = a;
    }
};

int main() {
    int N = 20;
    int a = 1, b = 2;

    DSU dsu(N);

    // a와 b 집합을 합침
    dsu.unite(a, b);

    // 원소 a와 b가 같은 집합에 포함되는지 여부를 확인
    if (dsu.get(a) == dsu.get(b)) 
}