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):

?

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45

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;     } }