코딩테스트에서 그래프 문제를 풀다 보면 "두 정점이 이미 연결되어 있나"를 계속 묻는 상황을 자주 만난다. 이걸 매번 BFS나 DFS로 확인하면 시간초과가 나기 쉽다. 이럴 때 꺼내는 자료구조가 Union-Find이고, 최소 신장 트리를 만드는 Kruskal 알고리즘의 핵심 부품이기도 하다. 예전에 알고리즘 글을 정리하면서 이 부분을 대충 넘어갔다가, 막상 실전에서 최적화 없이 짜서 TLE를 만나는 경우를 여러 번 봤다.
결론부터 말하자면, Union-Find는 find에 경로 압축을, union에 union by rank(또는 union by size)를 함께 적용하면 연산당 거의 상수 시간에 동작한다. Kruskal은 간선을 가중치 순으로 정렬한 뒤 사이클을 만들지 않는 간선만 골라 합치는 방식이고, 사이클 판정 역시 union이 실패하는지로 곧바로 알 수 있다. 아래에서 기본 구현, 두 가지 최적화, Kruskal 적용, 그리고 시간초과 함정을 차례로 살펴본다.
1. Union-Find(Disjoint Set)가 무엇이고 왜 쓰나
Union-Find는 서로소 집합(Disjoint Set), 즉 겹치지 않는 여러 집합을 관리하는 자료구조다. 지원하는 연산은 두 가지다.
- find(x): x가 속한 집합의 대표(루트)를 찾는다.
- union(a, b): a가 속한 집합과 b가 속한 집합을 하나로 합친다.
내부적으로는 각 원소가 부모를 가리키는 트리(정확히는 트리들의 숲)로 표현한다. 같은 루트를 가지면 같은 집합이다. 그래서 "두 정점이 연결되어 있나"는 find(a) == find(b) 한 줄로 판정한다. 여기서 rank는 트리 높이의 상한을 나타내는 값이고, 뒤에서 합치는 방향을 정하는 데 쓴다.
2. 기본 구현
부모 배열 parent를 자기 자신으로 초기화하는 것에서 시작한다. 최적화를 넣기 전의 뼈대는 이렇다.
class DisjointSet {
private final int[] parent;
private final int[] rank;
DisjointSet(int n) {
parent = new int[n];
rank = new int[n];
for (int i = 0; i < n; i++) parent[i] = i; // 처음엔 각자 자기 집합
}
int find(int x) {
while (parent[x] != x) x = parent[x]; // 루트까지 타고 올라감
return x;
}
boolean union(int a, int b) {
int ra = find(a), rb = find(b);
if (ra == rb) return false; // 이미 같은 집합
parent[rb] = ra;
return true;
}
}
이 상태로도 동작은 한다. 문제는 union을 아무 방향으로나 하다 보면 트리가 한쪽으로 길게 늘어져 사실상 연결 리스트가 되고, find가 O(n)까지 느려진다는 점이다. 여기서 두 가지 최적화가 필요하다.
3. 최적화 1: 경로 압축(Path Compression)
경로 압축은 find로 루트를 찾은 김에, 지나온 노드들의 부모를 곧바로 루트로 바꿔 트리를 납작하게 만드는 방법이다. 다음 조회부터는 한 번에 루트에 닿는다.
// 재귀 버전 (짧다)
int find(int x) {
if (parent[x] != x) parent[x] = find(parent[x]);
return parent[x];
}
재귀가 짧고 읽기 좋지만, 정점 수가 많고 트리가 깊게 만들어진 경우 재귀 깊이 때문에 StackOverflowError가 날 수 있다. 대회나 코딩테스트에서 정점이 수십만 개면 반복문 버전이 안전하다.
// 반복 버전 (스택 오버플로 방지)
int find(int x) {
int root = x;
while (parent[root] != root) root = parent[root];
while (parent[x] != root) { // 돌아오면서 부모를 루트로 교체
int next = parent[x];
parent[x] = root;
x = next;
}
return root;
}
경로 압축만 적용해도 연산당 평균 O(log n) 수준으로 좋아진다.
4. 최적화 2: Union by Rank
union by rank는 두 트리를 합칠 때 rank(높이 상한)가 작은 트리를 큰 트리 밑에 붙이는 방법이다. 낮은 트리를 높은 트리에 붙여야 전체 높이가 늘지 않는다. 두 rank가 같을 때만 붙인 쪽의 rank를 1 올린다.
boolean union(int a, int b) {
int ra = find(a), rb = find(b);
if (ra == rb) return false; // 이미 같은 집합 = 합칠 필요 없음
if (rank[ra] < rank[rb]) { int t = ra; ra = rb; rb = t; } // 큰 쪽을 ra로
parent[rb] = ra;
if (rank[ra] == rank[rb]) rank[ra]++;
return true;
}
경로 압축과 union by rank를 함께 쓰면 연산당 상각 시간 복잡도가 O(α(n))이 된다. 여기서 α(n)은 역 애커만 함수인데, 현실적인 모든 n(대략 10^600 이하)에서 4를 넘지 않아 사실상 상수로 취급한다. 이 결과는 1975년 Tarjan이 처음 증명했다.
rank 대신 집합 크기를 기준으로 큰 쪽에 붙이는 union by size도 있다. 성능은 사실상 같으니 둘 중 익숙한 걸 쓰면 된다. 중요한 건 경로 압축과 병합 기준을 둘 다 넣는 것이다.
5. Kruskal 최소 신장 트리 구현
최소 신장 트리(MST)는 모든 정점을 잇되 간선 가중치 합이 가장 작은 트리다. Kruskal은 간선을 가중치 오름차순으로 정렬한 뒤, 앞에서부터 하나씩 보면서 사이클을 만들지 않는 간선만 채택한다. 사이클 판정이 바로 Union-Find의 union이 false를 반환하는지로 끝난다.
import java.io.*;
import java.util.*;
public class Kruskal {
public static void main(String[] args) throws IOException {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
StringTokenizer st = new StringTokenizer(br.readLine());
int n = Integer.parseInt(st.nextToken()); // 정점 수
int m = Integer.parseInt(st.nextToken()); // 간선 수
int[][] edges = new int[m][3]; // {가중치, u, v}
for (int i = 0; i < m; i++) {
st = new StringTokenizer(br.readLine());
int u = Integer.parseInt(st.nextToken());
int v = Integer.parseInt(st.nextToken());
int w = Integer.parseInt(st.nextToken());
edges[i][0] = w; edges[i][1] = u; edges[i][2] = v;
}
Arrays.sort(edges, (e1, e2) -> Integer.compare(e1[0], e2[0]));
DisjointSet ds = new DisjointSet(n);
long total = 0;
int used = 0;
StringBuilder sb = new StringBuilder();
for (int[] e : edges) {
if (ds.union(e[1], e[2])) { // 사이클을 안 만들면 채택
total += e[0];
sb.append(e[1]).append(" - ").append(e[2])
.append(" (").append(e[0]).append(")\n");
if (++used == n - 1) break; // 간선 n-1개면 완성
}
}
System.out.println("MST 간선:");
System.out.print(sb);
System.out.println("총 가중치 = " + total);
}
}
DisjointSet 클래스는 앞서 만든 반복 버전 find와 union by rank union을 그대로 쓰면 된다.
입력 예시
정점 7개, 간선 11개. 각 줄은 u v w 형식이다.
7 11 0 1 7 0 3 5 1 2 8 1 3 9 1 4 7 2 4 5 3 4 15 3 5 6 4 5 8 4 6 9 5 6 11
실행 결과
MST 간선: 0 - 3 (5) 2 - 4 (5) 3 - 5 (6) 0 - 1 (7) 1 - 4 (7) 4 - 6 (9) 총 가중치 = 39
가중치 5짜리 간선 두 개를 먼저 먹고, 6, 7, 7, 9를 차례로 붙여 정점 7개를 잇는 간선 6개(n-1)로 트리가 완성된다. Kruskal의 전체 시간 복잡도는 간선 정렬이 지배해 O(E log E)다. Union-Find 연산은 간선마다 상각 상수라 정렬에 묻힌다.
6. 사이클 판정만 따로 쓰기
가중치 없이 "이 그래프에 사이클이 있나"만 볼 때도 같은 원리다. 간선을 순회하며 union이 실패하는 순간, 그 간선의 두 끝점은 이미 같은 집합이므로 사이클이 존재한다.
DisjointSet ds = new DisjointSet(n);
boolean hasCycle = false;
for (int[] e : edges) { // e = {u, v}
if (!ds.union(e[0], e[1])) { // 합치기 실패 = 이미 연결됨
hasCycle = true;
break;
}
}
System.out.println(hasCycle ? "사이클 있음" : "사이클 없음");
단, 이 판정은 무방향 그래프 기준이다. 방향 그래프의 사이클 판정은 DFS 방문 상태나 위상 정렬을 쓰는 편이 맞다. Union-Find는 연결 관계만 보지 방향을 구분하지 않는다.
7. 코딩테스트에서 자주 틀리는 TLE 함정
경로 압축이나 병합 기준을 빼먹는다
가장 흔한 실수다. find에 경로 압축이 없거나 union을 항상 같은 방향으로만 붙이면, 최악의 경우 트리가 한 줄로 늘어져 연산당 O(n)이 된다. 정점 10만 개 기준으로 이건 곧바로 시간초과로 이어진다. 둘 다 넣어야 상각 상수가 된다.
Scanner로 대용량 입력을 읽는다
간선이 수십만 줄이면 Scanner는 파싱 오버헤드로 느리다. BufferedReader와 StringTokenizer를 쓰면 입력 부분에서 나던 시간초과가 사라지는 경우가 많다. 위 예제도 그래서 BufferedReader로 짰다.
재귀 find로 스택 오버플로가 난다
재귀 find는 짧아서 좋지만, 압축 전 트리가 깊으면 재귀 깊이가 정점 수만큼 쌓여 StackOverflowError가 날 수 있다. 정점이 많은 문제에서는 4절의 반복 버전 find가 안전하다.
합계 자료형을 int로 둔다
MST 가중치 합은 간선 수와 가중치가 크면 int 범위를 넘긴다. 정답은 맞는데 오버플로로 음수가 나오는 어이없는 오답을 피하려면 누적 합은 long으로 받는다.
정렬 대상에 불필요한 박싱을 쓴다
간선을 int[][]로 두고 Arrays.sort에 Integer.compare를 쓰는 정도는 괜찮다. 다만 Integer[]나 무거운 객체 리스트로 정렬하면 박싱과 비교 비용이 커진다. 간선 개수가 많을수록 원시 배열 기반으로 짜는 편이 유리하다.
8. 마무리
Union-Find는 경로 압축과 union by rank(또는 size) 두 가지만 제대로 넣으면 연산당 사실상 상수 시간에 동작하는 값싼 자료구조다. Kruskal MST와 사이클 판정은 이 자료구조를 거의 그대로 얹은 응용이다. 코딩테스트에서 시간초과를 만난다면 최적화 누락, Scanner 입력, 재귀 깊이, 합계 자료형 이 네 가지부터 점검하면 대부분 잡힌다.