PS/USACO Gold

[USACO Gold] Disjoint Set, 분리 집합, Union-find

조금씩 차근차근 2025. 10. 7. 16:01

그래프에 간선을 추가하고 그래프의 두 정점이 연결되어 있는지 검사할 수 있게 해 주는 자료구조가 분리 집합(DSU, Union-Find) 이다.

USACO.guide는 알고리즘의 원리와 정의보단 어떻게 이 정의를 떠올릴 수 있는지에 집중한다.
알고리즘의 원리와 정의가 중요하지 않다는 뜻이 아니다.
실제 Union-find의 원리와 관련된 내용은 다음 링크를 참고 바란다.

 

CS Academy

 

csacademy.com

 

구현(Implementation)

import java.util.*;

public class DisjointSets {
    int[] parents;
    int[] sizes;

    public DisjointSets(int size) {
        parents = new int[size];
        sizes = new int[size];
        for (int i = 0; i < size; i++) {
            parents[i] = i;
            sizes[i] = 1;
        }
    }

    /** x가 속한 컴포넌트의 "대표" 노드를 반환 */
    public int find(int x) {
        return parents[x] == x ? x : (parents[x] = find(parents[x]));
    }

    /** 두 집합을 합치고, 연결성에 변화가 있었는지 여부를 반환 */
    public boolean unite(int x, int y) {
        int xRoot = find(x);
        int yRoot = find(y);
        if (xRoot == yRoot) { return false; }
        if (sizes[xRoot] < sizes[yRoot]) { return unite(yRoot, xRoot); }
        parents[yRoot] = xRoot;
        sizes[xRoot] += sizes[yRoot];
        return true;
    }

    /** x와 y가 같은 연결 컴포넌트에 속하는지 여부를 반환 */
    public boolean connected(int x, int y) { return find(x) == find(y); }
}

구현이 꽤 간단하기 때문에, 연결 요소를 계산할 때 DFS 대신 이 구조를 사용하는 편이 좋을 때가 많다.

문제 - Unionfind

다음 문제를 풀어보자.

 

Library Checker

 

judge.yosupo.jp

 

풀이 - 집중 문제(Focus Problem)

먼저 유니온-파인드를 모르는 상태로 생각해보자.
유니온 파인드를 쓰지 않으면, 그래프를 인접 리스트로 표현하고 플러드 필(flood fill) 로 연결 요소를 매번 계산해야 한다.
이 방법은 쿼리마다 연결 요소를 다시 세야 하므로 전체 시간 복잡도가 O(NQ)가 되어 너무 느리다.
이 문제는 "연결되어 있는가"라는 정보만 유지하면 풀 수 있기 때문에, 유니온 파인드를 사용한다.

위에서 구현한 유니온 파인드 자료구조로 그래프를 표현하면, 정점을 합치는 연산과 두 정점 ui,vi​ 가 같은 연결 요소에 속하는지 확인하는 연산을 모두 Amortized O(α(N)) 시간에 수행할 수 있다.

여기서 α는 매우 천천히 증가하는 함수(아커만 반전 함수)를 의미한다.

그 결과 전체 시간 복잡도는 O(Qα(N)) 로 줄어들며, 이는 큰 폭의 개선이므로 모든 테스트 케이스를 통과할 수 있게 된다.

import java.io.*;  

public class Main {  
    public static void main(String[] args) throws IOException {  
        int n, q;  
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));  
        PrintWriter pw = new PrintWriter(new OutputStreamWriter(System.out));  
        String s = br.readLine();  
        String[] split = s.split(" ");  
        n = Integer.parseInt(split[0]);  
        q = Integer.parseInt(split[1]);  
        DisjointSet ds = new DisjointSet(n);  
        for (int i = 0; i < q; i++) {  
            s = br.readLine();  
            split = s.split(" ");  
            int cmd = Integer.parseInt(split[0]);  
            int a = Integer.parseInt(split[1]);  
            int b = Integer.parseInt(split[2]);  
            if(cmd == 0){  
                ds.union(a, b);  
            }else{  
                if(ds.find(a) == ds.find(b)){  
                    pw.println("1");  
                }else{  
                    pw.println("0");  
                }  
            }  
        }  
        br.close();  
        pw.flush();  
        pw.close();  
    }  
}  

class DisjointSet {  
    int[] parent;  
    int[] rank;  

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

    public int find(int now) {  
        int p = parent[now];  
        if (p == now) return now;  
        return parent[p] = find(p);  
    }  

    public void union(int u, int v) {  
        u = find(u);  
        v = find(v);  
        if (u == v) return;  
        if (rank[u] < rank[v]) {  
            int temp = v;  
            v = u;  
            u = temp;  
        }  
        parent[v] = u;  
        if (rank[u] == rank[v]) rank[u]++;  
    }  
}

추천 문제