Union-Find 알고리즘
유니온 파인드(Union-Find)
시뮬레이션 문제를 풀다가 이런 상황을 만났다.
빙하 여러 개가 동시에 이동한다. 이동 후 서로 겹치거나 인접하면 하나로 병합된다.
병합 자체는 어렵지 않아 보였다. 문제는 이거였다.
- 0번과 3번이 겹쳤다
- 3번과 7번이 인접했다
- 0번과 5번이 인접했다 이 세 줄을 순서대로 처리하고 나면, 결국 0, 3, 5, 7이 전부 한 덩어리가 되어야 한다. 그런데 발견되는 순서는 제멋대로고, 사슬이 몇 개까지 이어질지도 모른다.
유니온 파인드는 정확히 이 상황을 위한 자료구조라고 한다.
1. 리스트로 하면 왜 안 되는가
먼저 자료구조 없이 해보자. 그룹을 리스트의 리스트로 관리한다.
groups = [[0], [1], [2], [3], [4], [5], [6], [7]]
union(0, 3)을 하려면:
- 0이 들어있는 리스트를 찾는다 → 전체 순회
- 3이 들어있는 리스트를 찾는다 → 전체 순회
- 두 리스트를 합친다
- 빈 리스트를 지우고 인덱스를 정리한다 여기까지도 번거로운데, 진짜 문제는 다음이다.
- 이미 합쳐진 걸 또 합치려는 경우를 매번 확인해야 한다
- 관계가 M개 들어오면 매번 선형 탐색이라 O(N × M)
- 코드가 금방 지저분해지고, 인덱스 관리에서 버그가 난다 "얘가 지금 어느 그룹이지?"를 빠르게 답할 수 있어야 한다는 게 본질이다.
2. 핵심 아이디어: 그룹을 트리로
유니온 파인드는 각 그룹을 트리 하나로 표현하고, 트리의 루트를 그룹 대표로 삼는다.
그리고 트리를 배열 하나로 표현한다.
parent[i] = i의 부모 (자기 자신이면 루트)
처음엔 전부 혼자다.
parent = list(range(n)) # [0, 1, 2, 3, 4, 5]
0 1 2 3 4 5 ← 각자가 자기 그룹의 루트
이제 필요한 연산은 두 개뿐이다.
find(a) = a가 속한 그룹의 대표(루트)를 반환
union(a, b) = a의 그룹과 b의 그룹을 하나로
이름 그대로 Find(찾기)와 Union(합치기)다. 그래서 Union-Find.
3. find — 루트까지 타고 올라가기
def find(a):
while parent[a] != a: # 내가 내 부모가 아니면 = 루트가 아니면
a = parent[a] # 부모로 한 칸 올라간다
return a
parent[a] == a인 지점이 루트다. 거기까지 올라가서 그 번호를 돌려준다.
4. union — 한쪽 루트를 다른 쪽 밑에 붙이기
def union(a, b):
ra, rb = find(a), find(b)
if ra != rb:
parent[rb] = ra
여기서 가장 중요한 건 a와 b를 직접 연결하는 게 아니라, 각자의 루트를 찾아서 연결한다는 점이다.
parent[b] = a # ❌ 그룹이 안 합쳐진다
parent[rb] = ra # ⭕ 그룹 전체가 합쳐진다
ra == rb면 이미 같은 그룹이므로 아무것도 하지 않는다. 덕분에 같은 union을 여러 번 호출해도 안전하다. 관계가 중복으로 들어오는 상황에서 이게 굉장히 편하다.
5. 손으로 따라가 보기
parent = [0, 1, 2, 3, 4, 5]
union(0, 3) → find(0)=0, find(3)=3, 서로 다르니 parent[3] = 0
parent = [0, 1, 2, 0, 4, 5]
0 1 2 4 5
|
3
union(5, 3) → find(5)=5, find(3)=0, 서로 다르니 parent[0] = 5
parent = [5, 1, 2, 0, 4, 5]
5 1 2 4
|
0
|
3
이제 find(3)을 해보면 3 → 0 → 5, 루트는 5. find(0)도 5, find(5)도 5.
0, 3, 5가 한 그룹이라는 사실이 자동으로 나온다. 우리가 알려준 건 "0-3", "5-3" 두 개뿐인데 말이다.
union(0, 3)을 또 호출하면? find(0)=5, find(3)=5로 같으므로 아무 일도 일어나지 않는다.
6. 최적화: 경로 압축
위 방식만으로도 정답은 나온다. 하지만 최악의 경우 트리가 일자로 길어진다.
0 → 1 → 2 → 3 → 4 → 5 → ...
이러면 find가 O(N)이 되어버린다. 그래서 보통 경로 압축(path compression) 을 같이 쓴다.
def find(a):
while parent[a] != a:
parent[a] = parent[parent[a]] # 올라가는 김에 링크를 당겨둔다
a = parent[a]
return a
추가된 한 줄이 하는 일은 "너 이제 할아버지가 부모야"라고 링크를 당겨두는 것이다.
5 5
| / \
0 find(3) 호출 후 → 0 3
|
3
한 번 훑을 때마다 트리가 납작해져서, 이후의 find가 훨씬 빨라진다.
처음 공부할 땐 이 줄을 빼고 보는 걸 추천한다. 동작이 훨씬 명확하게 보인다. 정답에는 영향이 없고, 속도만 달라진다.
유니온 바이 랭크(선택)
트리 높이를 기록해서 항상 낮은 트리를 높은 트리 밑에 붙이는 방식도 있다.
def union(a, b):
ra, rb = find(a), find(b)
if ra == rb:
return
if rank[ra] < rank[rb]:
ra, rb = rb, ra
parent[rb] = ra
if rank[ra] == rank[rb]:
rank[ra] += 1
경로 압축과 같이 쓰면 연산 하나당 사실상 상수 시간(정확히는 아커만 함수의 역함수 α(N))이 된다.
다만 코딩 테스트에서는 경로 압축만 써도 충분한 경우가 대부분이다. 랭크까지 챙기다가 버그 내는 것보다, 압축 한 줄만 넣고 확실하게 가는 편이 낫다.
7. 전체 구현 (복붙용)
parent = list(range(cnt))
def find(a):
while parent[a] != a:
parent[a] = parent[parent[a]]
a = parent[a]
return a
def union(a, b):
ra, rb = find(a), find(b)
if ra != rb:
parent[rb] = ra
8. 코테에서 언제 쓰는가
문제를 읽다가 아래 신호가 보이면 유니온 파인드를 의심하면 된다.
신호 1: "연결되어 있는지" 를 묻는다
- 두 노드가 같은 그룹인지 판정
- 네트워크가 전부 연결되었는지
- 친구의 친구까지 묶어서 세기
find(a) == find(b)한 줄이면 끝난다.
신호 2: 관계가 하나씩 추가된다
간선이 추가만 되고 삭제되지 않는 상황이 결정적인 힌트다. 유니온 파인드는 합치기만 가능하고 되돌리기(분리)는 못 하기 때문이다.
- 도로를 하나씩 건설하면서 상태를 추적
- 쿼리로 "a와 b를 연결" 이 계속 들어옴 반대로 간선이 삭제되는 문제라면 유니온 파인드를 그대로 쓸 수 없다. 이럴 땐 쿼리를 거꾸로 뒤집어서 삭제를 추가로 바꾸는 테크닉을 쓰곤 한다.
신호 3: 그룹 개수를 세야 한다
union이 실제로 합쳐질 때마다 카운트를 하나씩 줄이면 된다.
groups = n
for a, b in edges:
if union(a, b):
groups -= 1
신호 4: 최소 신장 트리 (크루스칼)
간선을 가중치 순으로 정렬한 뒤, 사이클을 만들지 않는 간선만 고른다. 이 "사이클을 만드는가" 판정이 정확히 유니온 파인드다.
edges.sort(key=lambda e: e[2])
total = 0
for a, b, w in edges:
if union(a, b):
total += w
신호 5: 여러 관계가 연쇄적으로 묶인다
이 글 서두의 빙하 문제가 여기 해당한다. A-B가 붙고 B-C가 붙으면 A-C도 같은 덩어리인데, 이 전이성을 직접 관리하려면 골치 아프다.
9. BFS/DFS와 뭐가 다른가
"연결 성분 구하기"는 BFS로도 된다. 실제로 그리드 문제에서는 BFS가 더 자연스러울 때가 많다.
핵심 차이는 동적이냐 정적이냐다. 관계가 중간에 계속 추가되는데 그때마다 BFS를 새로 돌리면 O(N × M)이 되지만, 유니온 파인드는 관계 하나당 거의 상수 시간이다.
그리고 유니온 파인드는 최단 경로를 알려주지 못한다. 트리 모양이 실제 연결 구조와 무관하게 압축되기 때문이다. "연결되어 있다"만 알 수 있고 "어떻게 연결되어 있다"는 모른다.