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

Example 1:

Input:
ConnectingGraph2(5)
query(1)
connect(1, 2)
query(1)
connect(2, 4)
query(1)
connect(1, 4)
query(1)
Output:[1,2,3,3]

Example 2:

Input:
ConnectingGraph2(6)
query(1)
query(2)
query(1)
query(5)
query(1)

Output:
[1,1,1,1,1]

思路：union find，老大哥root里面记录一下size的信息，connect的时候更新即可；注意不要在路径压缩的时候更新，因为只需要老大哥的信息，中间点的信息不需要；

public class ConnectingGraph2 {
    private class UnionFind {
        private int[] father;
        public int[] count;
        public UnionFind(int n) {
            this.father = new int[n + 1];
            this.count = new int[n + 1];
            for(int i = 0; i <= n; i++) {
                father[i] = i;
                count[i] = 1;
            }
        }

        public int find(int x) {
            int j = x;
            while(father[j] != j) {
                j = father[j];
            }
            // path compression;
            while(x != j) {
                int fx = father[x];
                father[x] = j;
                x = fx;
            }
            return j;
        }

        public void union(int a, int b) {
            int root_a = find(a);
            int root_b = find(b);
            if(root_a != root_b) {
                father[root_a] = root_b;
                count[root_b] += count[root_a];
            }
        }
    }
    /*
    * @param n: An integer
    */
    private UnionFind uf;
    public ConnectingGraph2(int n) {
        uf = new UnionFind(n);
    }

    /*
     * @param a: An integer
     * @param b: An integer
     * @return: nothing
     */
    public void connect(int a, int b) {
        uf.union(a, b);
    }

    /*
     * @param a: An integer
     * @return: An integer
     */
    public int query(int a) {
        int root_a = uf.find(a);
        return uf.count[root_a];
    }
}