파이썬으로 구현하는 시리즈

[파이썬으로 구현하는] 분리 집합(Disjoint Set)

kkigon 2025. 4. 5. 23:37

"파이썬으로 구현하는"

시리즈 1

분리 집합

(Disjoint Set)


오늘은 파이썬으로 구현하는 분리 집합의 표현 방법에 대해서 알아볼 것이다.

 

다른 말로 유니온-파인드 자료 구조(Union-Find Structure)라고도 하는 이 자료구조는, 집합의 묶음을 관리하는 구조이다. 여기에서 집합은 서로소 집합으로, 한 원소가 둘 이상의 집합에 속하는 경우가 없다.

 

아이디어는 다음과 같다.

 

집합 하나당 원소 하나가 대표값이 된다.

그래서, 같은 집합에 속해있으면 같은 대표값을 가지게 되는 것이다!

 

예를 들어, 집합이 {3, 1, 4}, {5} {6, 7, 8, 9}인 경우 다음과 같이 나타낼 수 있을 것이다.

 

꼭 이렇게 생겼다는 말은 아니다.

 

예를 들어서 나는 이 상태에서 {3, 1, 4}와 {6, 7, 8, 9} 집합을 합하고 싶다. 그러면 어떻게 하면 되겠는가?

바로 1번 노드가 6번 노드를 가리키게 하면 될 것이다!

 

이를 구현하기 위해, 우리는 각 원소들의 대표값을 저장하는 공간을 각 원소별로 가지고 있으면 될 것이다.

그래서, 우리는 처음에 parent라는 이름의 리스트를 만들어서 노드들의 대표값을 관리하도록 한다.

# 편의상 노드가 1, 2....N까지 있다고 하자
parent = [i for i in range(N+1)]

 

맨 처음에는 모든 원소의 대표값은 자기 자신이다.

현재 상태는 모두 다른 집합에 포함되어있다는 말이다.

 

우리는 두 개의 함수를 만들어 줄 것인데,

바로 각 원소의 대표값을 찾아주는 find() 함수와

 

두 원소가 같은 집합에 속하는지 판단하는 same() 함수,

 

두 함수를 합쳐주는 union() 함수이다.

 

이 두 개의 함수들은 앞으로 문제들을 풀면서 쿼리가 진행되면서, 적절히 사용해주면 될 것이다.

 

find() 함수의 경우, 대표값을 구하기 위해서는 각 원소의 대표값을 타고타고 찾아가주면서, 자기 자신을 가리키는 원소가 나오면 그것이 대표값이 되는 것이다(트리상에서 가장 위에 있는 노드).

이를 함수로 다음과 같이 구현할 수 있다.

 

def find(x):
    global parent
    if x == parent[x]:
        return x
    return find(parent[x])

 

same() 함수는 위에서 만든 find() 함수를 적절히 응용해서 만들 수 있다. 두 원소의 대표값이 같다면 같은 집합에 속하는 것이다.

 

def same(a, b):
	return find(a) == find(b)

 

union() 함수는 두 원소의 대표값들을 찾은 다음 하나의 대표값을 한 쪽에 연결시켜주면 될 것이다.

def union(a, b):
    global parent
    a = find(a)
    b = find(b)
    parent[b] = a

 

추가적인 기법으로, parent 리스트와 비슷하게 size 리스트를 만들어서 각 집합의 크기를 저장해준 다음, 크기를 비교해서 더 크기가 작은 집합을 더 큰 집합에 연결시키는 방법도 있다.

 


경로 압축(Path Compression)?

더 빠른 시간에 동작하게 하기 위해서, 우리는 경로 압축이라는 스킬을 적용해줄 수 있다.

다음 예를 보자.

5라는 원소에서 find 연산을 적용하면 우리는 그 부모 노드들을 차례차례 탐색해나가게 된다. 만약 이 체인이 길다면 탐색에 시간이 조금 걸릴 것이다.

근데 만약, 이 부모들을 그냥 바로 대표값으로 직속 연결해버리면 나중에 찾을 일이 있을 때 시간이 더 적게 걸리지 않을까?

 

실제로, 이 함수를 이용하면 유니온-파인드 연산의 분할 상환에 따른 시간 복잡도가 ((O(\alpha(n))))이 된다고 알려져있다. 

여기서 ((\alpha(n)))은 애커만 함수(ACKERMANN FUNC.)의 역함수로, 증가 속도가 매우매우 느려 상수라고 봐도 무방하다!

즉, 우리는 find() 연산을 거의 ((O(1)))안에 해결할 수 있다!

 

이 경로 압축 연산은 다음과 같이 find() 함수를 조금 바꾸어서 만들 수 있다.

 

def find(x):
    global parent
    if x == parent[x]:
        return x
    parent[x] = find(parent[x])
    return parent[x]

 

 

 

 

그럼 이제 최종적인 구현 코드들을 정리해보자.

parent = [i for i in range(N+1)]

def find(x):
    global parent
    if x == parent[x]:
        return x
    parent[x] = find(parent[x])
    return parent[x]

def union(a, b):
    global parent
    a = find(a)
    b = find(b)
    parent[b] = a

 

 

재귀 안쓰는 버전(25.08.04)

문제풀다보니까 아무리 path compression을 해도 최악의 경우 재귀 한도에 걸려서 메모리와 시간을 많이 잡아먹는다는 것을 알게 되었다.

그래서, 새로운 재귀 안쓰는 버전의 find를 소개한다.

def find(x):
    global parent
    while x != parent[x]:
        parent[x] = parent[parent[x]]
        x = parent[x]
    return x

 

첨언(25.09.05)

생각해보니까 위 코드는 효과적으로 경로 압축을 하지는 않는다.

그냥 부모 노드쪽으로 1칸 이동시켜줄 뿐...

결국 재귀가 답인건가...?


 

참고로 재귀 버전 대비 시간 차이는...

 

 

엄청나다! 거의 5배차이로 빠르다.

 

 

대표적인 연습문제로는

 

집합의 표현

https://www.acmicpc.net/problem/1717

 

여행 가자

https://www.acmicpc.net/problem/1976

 

거짓말

https://www.acmicpc.net/problem/1043

 

피리 부는 사나이

https://www.acmicpc.net/problem/16724

 

등이 있다.

 

 

 

 

참고: <알고리즘 트레이닝 2판> - 안티 라크소넨