Union-Find: 상호 배타적 집합을 표현하는 Java 코드

Union-Find

공통 원소가 없는 상호 배타적인 부분 집합들로 나눠진 원소들을 저장, 조작하는 자료구조

Java code

public class UnionFind {
    int[] root;
    int[] rank;
    
    public UnionFind(int N){
        root = new int[N];
        rank = new int[N];
        for(int i=0; i<N; i++){
            root[i] = i;
            rank[i] = 1;
        }
    }
    
    public int find(int a){
        if(a == root[a]) return a;
        return root[a] = find(root[a]);
    }
    
    public int union(int a, int b){
        if(a == -1 && b == -1) return -1;
        if(a == -1) return find(b);
        if(b == -1) return find(a);
        a = find(a);
        b = find(b);
        if(a == b) return a;
        if(rank[a] > rank[b]){ //항상 높이가 작은걸 큰것에 붙임
            int temp = a;
            a = b;
            b = temp;
        }
        root[a] = b;
        if(rank[a] == rank[b]) rank[b]++;
        return b;
    }
}