[Python] 분리 집합, Union-Find에 대해서 알아보자!

분리집합(Union-Find)은 그래프 알고리즘 중 하나로 코딩 테스트 킬러 문제로 주로 나오는 알고리즘입니다. 

 

서로소 집합

일단 유니온 파인드 알고리즘을 알기 전 disjoint sets을 먼저 이해해야 하는데요, 쉽게 설명하면 공통 원소가 없는 두 집합을 뜻합니다. 

아래 그림을 보시면 왼쪽은 서로소 집합을 나타내고 오른쪽은 공통된 7이라는 원소 때문에 서로소 관계가 깨지게 됩니다. 따라서 왼쪽은 서로소 집합, 오른쪽은 서로소 집합이 아니게 되는 것이죠. 

https://www.math-only-math.com/disjoint-of-sets-using-Venn-diagram.html

Union-Find는 서로소 집합 자료구조를 만들 수 있습니다. 집합에서 노드들을 합치고(Union), 부모를 찾아(Find) 서로소 집합을 찾아내는 알고리즘인 것이죠. 한마디로 DSU(Disjoint Sets Union) 자료구조를 만들어 내는 과정이라고 볼 수 있습니다. 기본적으로 Union-Find는 트리 자료구조에서 작동합니다. 

 

Union 연산

Union은 2개의 set을 하나의 set으로 합쳐주는 연산입니다.

2개의 서로소 집합을 Union 연산을 통해 하나로 합치는 과정(합집합)을 거치면 다음과 같은 트리가 완성됩니다.

기본적인 룰은 각 집합의 root node를 찾아 더 높은 인덱스를 가진 쪽이 작은 쪽을 흡수하는 성질을 가지고 있습니다.

이를 파이썬에서 동작 시켜보면 다음과 같은 함수로 구현할 수 있습니다.

parent = [i for i in range(8)]

def union(a, b):
    a = find(a)
    b = find(b)
    if a < b:
        parent[a] = b
    else:
        parent[b] = a

 union 함수를 살펴보면 find 함수가 나오게 되는데, 이 find(x) 함수는 각 x 인자의 루트 노드를 찾는 함수입니다. 위 그림의 예시로 각 루트노드의 값이 2와 6이었습니다. 아래 표를 통해 위 그림과 같은 상황에서 union 함수를 작동했을 때 어떻게 되는지 봅시다.

Union 연산 전
index(노드번호)	1	2	3	4	5	6	7
부모 노드	2	2	2	5	6	6	5

union 연산 전 배열의 상태를 나타낸 그림입니다. 배열의 index가 노드 번호를 나타내고 있고, 배열 안 값이 해당 노드의 부모 노드를 가르키고 있습니다. 

Union 연산 이후
index(노드번호)	1	2	3	4	5	6	7
부모 노드	2	6	2	5	6	6	5

union 연산 이후 2번 노드의 부모 노드가 6으로 변경됨으로 Union 연산을 마쳤습니다. 

 

Find 연산

그럼 어떠한 과정으로 Union 연산에서 루트 노드를 찾았는지 확인 해 봅시다. 

Find 연산은 어떤 노드가 주어질 때 이 요소의 루트 노드를 찾아주는 연산입니다. 즉 원소가 속한 집합을 알려주는 연산과도 같죠. 

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

for i in range(n):
	find(i)

find 함수는 각 노드의 루트 노드를 재귀로 찾아가는 함수입니다. for문을 통해 find 연산을 n번만큼 해주고 나면 배열은 다음 표처럼 변할 것입니다.

Union 연산 이후
index(노드번호)	1	2	3	4	5	6	7
부모 노드	6	6	6	6	6	6	6

모든 노드가 6을 가르키고 있으므로 하나의 서로소 집합인 것을 확인할 수 있습니다. 

 

백준 1717번: 집합의 표현

유니온 파인드 문제를 간단하게 백준에서 풀어보겠습니다. 

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

1717번: 집합의 표현

0 또는 1 입력을 통해 집합을 합치거나 확인 하는 연산을 해주기만 하면 됩니다. 합치는 연산은 Union, 찾는 연산은 find이기 때문에 0이 맨 앞에 들어올 경우 union 연산을 1이 들어올 경우엔 a, b 둘 다 find 연산을 통해 두 수의 루트 노드가 같은지 즉, 같은 서로소 집합 내에 있는지 확인해보면 됩니다. 

import sys
input = sys.stdin.readline
n, m = list(map(int, input().split()))
parent = [i for i in range(n + 1)]

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

def union(a, b):
    a = find(a)
    b = find(b)
    if a < b:
        parent[a] = b
    else:
        parent[b] = a

for i in range(m):
    x, a, b = list(map(int, input().split()))
    if not x:
        union(a, b)
    else:
        if find(a) == find(b):
            print("YES")
        else:
            print("NO")