Union-Find — 두 최적화가 O(N) 을 상수로 만든다

그래프 문제에는 “두 원소가 같은 그룹인가"를 반복해서 묻는 유형이 있다. 친구의 친구까지 이으면 그룹이 몇 개인지, 간선을 하나 추가하면 사이클이 생기는지. 그래프 토픽을 정리하면서 이 질문의 표준 도구인 Union-Find(서로소 집합, Disjoint Set)를 최적화 없는 버전부터 구현해 봤는데, 최적화 두 개가 왜 반드시 붙는지 확인할 수 있었다. 없으면 find 가 O(N) 까지 늘어지고, 둘 다 붙이면 사실상 상수가 된다.

find 와 union

각 그룹에 대표(root)를 하나 두고, 모든 원소가 parent 를 따라가면 자기 그룹의 root 에 도달하게 만든다. 두 원소의 root 가 같으면 같은 그룹이다. find(x) 는 x 에서 parent 를 따라가 root 를 반환하고, union(x, y) 는 두 root 를 찾아 한쪽을 다른 쪽에 붙인다.

class UnionFindNaive:
    def __init__(self, n):
        self.parent = list(range(n))

    def find(self, x):
        while self.parent[x] != x:
            x = self.parent[x]
        return x

    def union(self, x, y):
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False
        self.parent[root_y] = root_x
        return True

이 구현도 정답은 낸다. 문제는 트리의 모양이다. union(1, 0), union(2, 1), union(3, 2), union(4, 3) 순서로 호출하면 매번 새 원소가 기존 트리의 root 를 자식으로 삼아, 0 → 1 → 2 → 3 → 4 로 한 줄로 이어진 사슬(chain) 트리가 된다. 이제 find(0) 은 root 까지 네 단계를 올라간다. 원소가 N 개면 find 최악이 O(N) 이고, union 도 내부에서 find 를 부르므로 N 번 반복하면 O(N²) 다.

flowchart BT
    c0["0"] --> c1["1"] --> c2["2"] --> c3["3"] --> c4["4"]

Path Compression

첫 번째 최적화는 find 안에 있다. root 를 찾으러 올라가면서 경로에 있던 노드들의 parent 를 root 로 직접 바꾼다. 다음 find 부터는 그 노드들 모두 한 단계면 root 에 도달한다.

def find(self, x):
    if self.parent[x] != x:
        self.parent[x] = self.find(self.parent[x])
    return self.parent[x]

재귀가 root 까지 올라갔다가 되돌아오면서 경로의 모든 노드를 root 에 직결한다. 위의 사슬 트리에서 find(0) 을 한 번 부르면 트리가 이렇게 바뀐다.

flowchart BT
    f0["0"] --> f4["4"]
    f1["1"] --> f4
    f2["2"] --> f4
    f3["3"] --> f4

자주 조회되는 트리일수록 빠르게 평탄해진다.

Union by Rank

union 쪽 최적화는 트리를 붙이는 방향을 정한다. 아무 쪽이나 붙이면 사슬이 생길 수 있으니, 높이가 낮은 트리를 높은 트리에 붙인다. 전체 높이는 두 트리의 높이가 같을 때만 1 늘어난다.

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x, y):
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False
        if self.rank[root_x] < self.rank[root_y]:
            root_x, root_y = root_y, root_x
        self.parent[root_y] = root_x
        if self.rank[root_x] == self.rank[root_y]:
            self.rank[root_x] += 1
        return True

rank 는 압축 전 트리 높이의 상한이다. path compression 이 실제 높이를 줄여도 rank 는 갱신하지 않는데, 병합 방향을 정하는 기준으로는 그대로 충분하다. rank 대신 그룹 크기(size)를 기준으로 작은 트리를 큰 트리에 붙여도 같은 보장이 성립하고, 그룹 크기 조회가 필요하면 size 쪽이 낫다.

두 최적화를 함께 쓰면 find 와 union 은 상각(Amortized) 기준 O(α(N)) 이다. α 는 애커만 함수의 역함수(Inverse Ackermann)로 현실적인 모든 N 에서 4 를 넘지 않아, 사실상 상수로 취급한다.

연결 요소 카운트

Number of Provinces(LeetCode 547)는 인접 행렬로 친구 관계가 주어지고 그룹 개수를 묻는다. 연결 요소(Connected Component) 수를 N 으로 시작해 union 이 성공할 때마다 1 줄이면, 순회가 끝난 시점의 값이 답이다.

def count_provinces(is_connected):
    n = len(is_connected)
    uf = UnionFind(n)
    count = n
    for i in range(n):
        for j in range(i + 1, n):
            if is_connected[i][j] and uf.union(i, j):
                count -= 1
    return count

행렬이 대칭이라 상삼각(i < j)만 순회하면 된다. 이미 같은 그룹이면 union 이 False 를 돌려주므로 count 를 잘못 줄이는 일이 없다.

이 반환값은 사이클 검출로도 이어진다. 간선을 하나씩 추가하다가 union 이 False 를 돌려주는 순간, 두 끝점이 이미 연결되어 있었다는 뜻이므로 그 간선이 사이클을 만든다. Redundant Connection(LeetCode 684)이 이 패턴 그대로다.

선택 기준

연결 요소를 다루는 도구로는 BFS/DFS 도 있다.

항목	BFS/DFS	Union-Find
그래프	고정, 한 번 순회	간선이 점진 추가
얻는 것	경로, 방문 순서	소속, 그룹 수, 사이클 판정
쿼리	순회 후 일괄	추가와 조회 interleave

BFS/DFS 가 적합한 경우:

그래프가 고정되어 있고 한 번의 순회로 답이 나온다
연결 요소의 내용물(경로, 방문 순서)이 필요하다
격자 탐색처럼 이웃이 좌표로 계산된다

Union-Find 가 적합한 경우:

간선이 점진적으로 추가되면서 중간중간 연결 여부를 물어본다
간선 추가가 사이클을 만드는지 즉시 판정해야 한다
Kruskal MST 처럼 “이 간선을 채택해도 되는가"를 반복 판정한다

같은 그룹인지 묻는 문제는 root 대표 하나로 판정이 끝나고, 그 판정을 빠르게 유지하는 건 두 최적화의 몫이다. 최적화 없는 버전을 먼저 구현해 보고 나서야 이 둘이 왜 항상 함께 붙는지 이해하게 됐다.

find 와 union#

Path Compression#

Union by Rank#

연결 요소 카운트#

선택 기준#

find 와 union

Path Compression

Union by Rank

연결 요소 카운트

선택 기준