🚶‍♂️Road to 코테(14) - Union-Find

Union-Find(=Disjoint Set Union, DSU)는 “서로소 집합”을 효율적으로 관리하면서, 두 원소가 같은 집합인지 확인(Find)하고 두 집합을 합치는(Union) 연산을 빠르게 처리하는 자료구조/알고리즘이다. 경로 압축(Path Compression)과 랭크/사이즈 기반 합치기(Union by Rank/Size)를 같이 쓰면, 연산이 거의 상수 시간에 가깝게 동작한다.

Union-Find가 필요한 순간

그래프에서 “사이클이 생기는지” 확인하거나, 최소 신장 트리(MST)인 Kruskal 알고리즘을 구현할 때 자주 등장한다. 특히 Kruskal에서는 간선을 가중치 순으로 보며, 두 정점이 이미 같은 집합이면 그 간선을 추가할 때 사이클이 생기므로 제외하는 방식으로 Union-Find를 사용한다.

핵심 연산 2가지

• Find(x): x가 속한 집합의 대표(루트)를 찾는다.
• Union(a, b): a가 속한 집합과 b가 속한 집합을 합친다(이미 같으면 no-op이다).

여기서 DSU는 parent 배열로 트리(포레스트)를 표현하는 방식이 일반적이고, Find는 트리의 루트까지 따라 올라가는 과정이다.

최적화 2가지 (중요)

• Path Compression: Find를 수행할 때, 거쳐간 노드들의 parent를 루트로 직접 연결해 다음 Find를 빠르게 만든다.
• Union by Rank: 두 트리를 합칠 때 랭크(대략 높이)가 낮은 쪽을 높은 쪽 밑에 붙여 트리가 한쪽으로 길어지는 것을 방지한다.

Java 구현 (Path Compression + Union by Rank)

아래 코드는 코딩테스트에서 바로 가져다 쓰기 좋은 형태(0..n-1 정점)로 작성한다.

public class UnionFind {
    private final int[] parent;
    private final int[] rank; // 트리의 대략적 높이다.

    public UnionFind(int n) {
        parent = new int[n];
        rank = new int[n];
        for (int i = 0; i < n; i++) parent[i] = i;
    }

    // Find with path compression이다.
    public int find(int x) {
        if (parent[x] == x) return x;
        parent[x] = find(parent[x]); // path compression이다.
        return parent[x];
    }

    // Union by rank이다.
    public boolean union(int a, int b) {
        int ra = find(a);
        int rb = find(b);
        if (ra == rb) return false;

        if (rank[ra] < rank[rb]) {
            parent[ra] = rb;
        } else if (rank[ra] > rank[rb]) {
            parent[rb] = ra;
        } else {
            parent[rb] = ra;
            rank[ra]++;
        }
        return true;
    }

    public boolean isSameSet(int a, int b) {
        return find(a) == find(b);
    }
}

예시: 사이클 판별(무방향 그래프)

무방향 그래프에서 간선 (u, v)를 추가하려는데 u와 v가 이미 같은 집합이면 그 간선은 사이클을 만든다.

// edges: int[][] edges = { {0,1}, {1,2}, {2,0} ... }이다.
UnionFind uf = new UnionFind(n);

boolean hasCycle = false;
for (int[] e : edges) {
    int u = e[0], v = e;
    if (uf.isSameSet(u, v)) { // 이미 연결됨 => 추가하면 사이클이다.
        hasCycle = true;
        break;
    }
    uf.union(u, v);
}