Union Find

UnionFind.java

class UnionFind {
    public int[] rank;
    public int[] parent;
    public int n;
    
    public UnionFind(int n){
        this.n = n;
        rank = new int[n];
        parent = new int[n];
        for(int i = 0 ; i < n ; i++)
            parent[i] = i;
    }
    
    public int find(int node) {
        if(parent[node] != node)
            parent[node] = find(parent[node]);
        return parent[node];    
    }
    
    public void union(int a, int b) {
        int root_a = find(a);
        int root_b = find(b);
        
        if(root_a == root_b)
            return;
        
        if(rank[root_a] < rank[root_b]) 
            parent[root_a] = root_b;
        
        else if(rank[root_a] > rank[root_b])
            parent[root_b] = root_a;
        
        else {
            parent[root_b] = root_a;
            rank[root_a]++; 
        }
    }
}

public class main {
    public static void main(String[] argv) {
        UnionFind rel = new UnionFind(5);
        
        rel.union(0, 2);
        rel.union(4, 2);
        rel.union(3, 1);
        
        for(int i : rel.rank)
            System.out.print(i + " ");
        
        for(int i : rel.parent)
            System.out.print(i + " ");
        
        System.out.println(rel.find(4) == rel.find(0));
        System.out.println(rel.find(1) == rel.find(0));
    }
}