[Algorithm] Union-Find

1. Union-Find란?

 

여러 개의 원소가 있고, 여러 개의 집합이 있다고 가정하자. 특정 원소가 어느 집합에 속해있는지 확인하고, 특정 집합을 합쳐야 할때 해당 알고리즘을 사용한다. 집합에 속하는 것을 한 원소를 대표로 하여 자식 노드로 이어줌으로써 집합을 구현한다. 

 

1) Union

두 집합을 합치는 연산이다. 두 집합에 속하는 자식 노드들을 모두 찾은 후, 한 집합의 부모 노드를 다른 집합의 자식 노드로 하면 모든 원소가 하나의 부모 노드로 이어지게 된다. 코드로 표현하자면, 부분집합 x의 조상(루트)을 부모로 하고 y를 그 밑에 붙인다고 생각하면, uf(find[x]) = uf(find[y]) 가 된다.

 

2) Find

집합에 속한 원소를 찾는 연산이다. 부모 노드에 대해 연결된 자식 노드를 찾고, 그 자식 노드의 자식 노드를 더 이상 없을 때까지 계속 찾는 과정으로 이루어진다. 

그림으로 보면 다음과 같다. 7개의 노드가 있다고 가정하자. 

초기 상태

 

uf 배열은 자신의 부모 노드를 원소로 가지는 배열이다. 처음에는 부모, 자식 관계가 없으니 -1로 초기화한다.

(0, 2), (4, 5) 두 개의 부분집합이 생겼다고 가정하자. 그 때 배열의 변화는 다음과 같다. 

첫 번째 부분집합 생성

누가 부모가 될지는 본인 마음이다. 어차피 소속 여부를 따지는 것이기 때문에 순서는 상관 없다.

그렇게 되면 2는 0의 부모, 5는 4의 부모가 된다. 이 때 find(0) 연산을 수행하면 2를 반호나할 것이고, find(5) 연산을 수행하면 4를 반환할 것이다. 이로써 각각이 같은 집합에 있다는 걸 알 수 있다. 

 

이제 두 집합을 합쳐보자. 위에서 말한 것과 같이, 하나의 부모 노드를 다른 자식 노드의 자식 노드로 만든다. 

여기서는 2를 조상으로 하고 나머지를 모두 자식으로 만들겠다.

 

union 연산 이후

2의 자식인 0을 5의 부모로 만들면 집합이 합쳐진 효과가 일어난다.

 

 

그러나 위와 같이 구현하면, 집합이 하나밖에 없을 경우 부모 노드에 대해 자식 노드가 일렬로 이어져 있으므로 순회하는데 O(n)의 시간이 소요된다. 이를 해결하기 위해 find를 호출할 때 조상을 찾아내면, 현재 값의 부모 값을 조상으로 바꾸는 것이다. 다시 말해서, 줄줄이 연결되어있던 자식 노드들이 여러 개의 선을 이용해 하나로 이어지는 것이다. 

 

위의 예제를 그대로 차용하자면, find(4)를 호출하면 5을 찾고, 재귀적인 구조이므로 find(5)가 호출되어 0을 반환한다. 다시 find(0)이 호출되어 2를 반환하고, 2는 조상이 없으므로 find 연산을 종료한다.

 

이 연산 과정에서 2를 제외한 모두는 2가 조상임을 알게 된다. 따라서 재귀적인 연산이 반환하는 최종 값, 즉 조상을 부모 값으로 한다. 그림으로 나타내면 다음과 같다.

 

코드로 표현하자면, uf[x] = find(x)였던 기존의 재귀적 구조가 uf[x] = find(uf[x])로 바뀌게 되어 자신의 조상을 계속 바꿀 수 있는 구조가 된다.

2. 구현

 

uf = [-1 for _ in range(n)]
def find(uf, x):
    if uf[x] == -1:
        return x
    tmp = find(uf, uf[x])
    uf[x] = tmp
    return uf[x]

def union(uf, x, y):
    X, Y = find(uf, x), find(uf, y)
    uf[X] = Y
    return uf