💡파이썬 Disjoint-Set(Union-Find) 이론 및 예제

알고리즘 코딩 테스트를 준비하다 보면 다양한 자료구조와 알고리즘을 만나게 됩니다.

그 중에서도 Disjoint-Set(서로소 집합)을 활용한 Union-Find는 중상난이도 문제에서 자주 출제되는 유형 중 하나입니다.

이번 글에서는 Disjoint-Set, Union-Find의 개념과 파이썬으로 구현하는 방법에 대해 알아보겠습니다.

Disjoint-Set(Union-Find)이란?

서로소 집합

• 단순히 disjoint-set은 서로 다른 두 집합(서로소 집합)을 뜻함
• 서로 다른 집합을 모아 하나의 집합으로 만드는 Union-find 알고리즘에 사용됨
• 사용되는 case
  • 그래프에서 서로 다른 정점을 연결하는 최소 신장 트리를 구하는 문제
  • 서로 다른 집합을 합치는 문제 등
• 수행 연산
  • Union: 두 개의 집합을 하나로 합치는 연산
  • Find: 주어진 원소가 속하는 집합의 루트 노드를 찾는 연산
• 트리(Tree) 자료구조를 이용하여 구현됨

파이썬으로 Union-Find 구현하기

단순히 말 그대로의 disjoint-set은 파이썬의 set 자료구조를 사용하면 다음과 같이 사용할 수 있습니다.

a_set = {1,2,3}
b_set = {4,5}
a_set.isdisjoint(b_set) # True
a_set | b_set # {1,2,3,4,5}

하지만 위와 같이 구현하면, 서로소 집합 갯수만큼의 변수 또는 서로소 집합 갯수 크기의 집합 활용이 필요하므로 공간복잡도가 불필요하게 증가합니다.

또한 서로소 집합 갯수가 많아질수록 각 집합의 disjoint 여부를 체크하다보면 시간복잡도도 크게 증가합니다.

특정 두 원소가 한 집합 안에 속하는지 확인하는 것 또한 어려움이 있습니다.

 

따라서 다음과 같이 하나의 집합으로, 트리(Tree)구조를 활용하여 구현합니다.

여기서 집합의 각 인덱스는 해당 원소의 노드 번호,

집합의 각 원소의 값은 부모 노드의 인덱스를 뜻합니다. (-1일 경우 자기 자신이 부모 노드. 즉, 트리의 root)

class DisjointSet:
    def __init__(self, size):
        self.data = [-1] * size # 초깃값 -1

    def find(self, x):
		# recursion으로도 구현해 되나, 가능하면 반복문으로 구현 (recursion limit error때문)
        while True:
            parent = self.data[x]
            if parent < 0:
                return x
            x = parent

    def union(self, x, y):
        root_x, root_y = self.find(x), self.find(y)
        if root_x != root_y:
            self.data[root_y] = root_x

 

웬만한 문제는 위와 같이 구현하면 다 풀리지만, 위와 같이 구현했음에도 time limit exceed되는 경우가 있습니다.

편향 Tree

위와 같이, 데이터가 편향된 Tree들이 주어졌을 경우, 위 구현에서 find의 시간복잡도는 O(N)이 되기 때문입니다.

그런 경우 각 노드의 tree에서의 rank 개념을 도입하여, 다음과 같이 최적화해주면 됩니다.

class DisjointSet:
    def __init__(self, size):
        self.data = [-1] * size
        self.rank = [0]

    def find(self, x):
        while True:
            parent = self.data[x]
            if parent < 0:
                return x
            x = parent

    def union(self, x, y):
        root_x = self.find(x)
        root_y = self.find(y)

        if root_x == root_y:
            return

        if self.rank[root_x] > self.rank[root_y]:
            self.data[root_y] = root_x
        elif self.rank[root_x] < self.rank[root_y]:
            self.data[root_x] = root_y
        else:
            self.data[root_y] = root_x
            self.rank[root_x] += 1

 

예제

boj의 친구 네트워크가 Union-Find를 활용한 유명한 문제 중 하나입니다.

풀이 방법은 다양하겠지만, 큐 활용을 잘 생각해보면 간단히 풀 수 있는(40줄 이내) 문제이니 풀어보시고 다른 사람들의 풀이도 꼭 참고해보시길 바랍니다.