(C++) Union-find 구현

명세

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

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

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

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

구현

vector<int> p(n + 1), size(n + 1, 1);
for (int i = 1; i <= n; i++) p[i] = i;

int get(int a, vector<int>& p) {
    if (p[a] == a) return a;
    return p[a] = get(p[a], p);
}
void unite(int a, int b, vector<int>& p, vector<int>& size) {
a = get(a, p);
b = get(b, p);
if (a == b) return;

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

// a와 b 집합을 합침
unite(a, b, p, size)

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