Lintcode 590. 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]




Code (Java):

public class ConnectingGraph2 {
    private int[] parents;
    private int[] counts;

    public ConnectingGraph2(int n) {
        parents = new int[n + 1];
        counts = new int[n + 1];

        for (int i = 0; i <= n; i++) {
            parents[i] = i;
            counts[i] = 1;
        }
    }

    public void connect(int a, int b) {
        int rootA = find(a);
        int rootB = find(b);

        if (rootA != rootB) {
            parents[rootB] = rootA;
            counts[rootA] += counts[rootB];
        }
    }

    public int query(int a) {
        int rootA = find(a);
        return counts[rootA];
    }

    private int find(int a) {
        int x = a;

        while (x != parents[x]) {
            x = parents[x];
        }

        while (a != x) {
            int temp = parents[a];
            parents[a] = x;
            a = temp;
        }

        return x;
    }
}