[LintCode] Connecting Graph II

Given n nodes in a graph labeled from 1 to n. There is no edges in the graph at beginning.

You need to support the following method:
1. connect(a, b), an edge to connect node a and node b
2. query(a), Returns the number of connected component nodes which include node a.

 

Example

5 // n = 5
query(1) return 1
connect(1, 2)
query(1) return 2
connect(2, 4)
query(1) return 3
connect(1, 4)
query(1) return 3

class UnionFind {
    private int[] father = null;
    private int[] numsInUnion = null;
    
    public UnionFind(int n){
        father = new int[n];
        numsInUnion = new int[n];
        for(int i = 0; i < n; i++){
            father[i] = i;
            numsInUnion[i] = 1;
        }
    }
    public int find(int x){
        if(father[x] == x){
            return x;
        }
        return father[x] = find(father[x]);
    }
    public void connect(int a, int b){
        int root_a = find(a);
        int root_b = find(b);
        if(root_a != root_b){
            father[root_a] = root_b;
            numsInUnion[root_b] += numsInUnion[root_a];
        }
    }
    public int queryNumInUnion(int a){
        return numsInUnion[find(a)];
    }
}
public class ConnectingGraph2 {
    private UnionFind uf = null;
    public ConnectingGraph2(int n) {
        uf = new UnionFind(n);
    }

    public void connect(int a, int b) {
        uf.connect(a - 1, b - 1);
    }
        
    public int query(int a) {
        return uf.queryNumInUnion(a - 1);
    }
}