【Lintcode】590. Connecting Graph II

题目地址：

https://www.lintcode.com/problem/connecting-graph-ii/description

设计一个有向图数据结构，初始化$n$个顶点，一开始没有边，每次可以执行两个操作：
1、给定两个顶点$a$和$b$，建立一条边；
2、给定一个顶点$a$，返回$a$所在连通块的顶点个数。

思路是并查集。因为要返回某个顶点所在连通块个数，我们可以维护一个size数组，对于森林里的每个树根v，size[v]就表示这棵树的节点个数。代码如下：

import java.util.Arrays;

public class ConnectingGraph2 {
    
    private int[] parent, size;
    // 给定一个顶点，返回其所在树的树根
    private int find(int x) {
        if (x != parent[x]) {
            parent[x] = find(parent[x]);
        }
        
        return parent[x];
    }
    
    /*
     * @param n: An integer
     */
    public ConnectingGraph2(int n) {
        // do intialization if necessary
        parent = new int[n];
        for (int i = 0; i < parent.length; i++) {
            parent[i] = i;
        }
        
        size = new int[n];
        Arrays.fill(size, 1);
    }
    
    /*
     * @param a: An integer
     * @param b: An integer
     * @return: nothing
     */
    public void connect(int a, int b) {
        // write your code here
        int pa  = find(a - 1), pb = find(b - 1);
        if (pa != pb) {
            size[pb] += size[pa];
            parent[pa] = pb;
        }
    }
    
    /*
     * @param a: An integer
     * @return: An integer
     */
    public int query(int a) {
        // write your code here
        return size[find(a - 1)];
    }
}

两个操作的时间复杂度都是$O({\log }^{\ast }n)$，空间$O(n)$（size数组的空间）。