I'm currently enrolled in the Princeton Algorithms course (Part 1) and it talks about an improvement to the quick-union algorithm by maintaining an extra array sz[i] to count the number of objects in the tree rooted i, but it doesn't show how to do that.
Where and how is that counter supposed to be implemented? I've tried doing it in the root method, but I realized it wouldn't count the children of a given object.
This is the unaltered code given in the course:
public class QuickUnionUF {
private int[] id;
public QuickUnionUF(int N) {
id = new int[N];
for (int i = 0; i < N; i++) id[i] = i;
}
private int root(int i) {
while (i != id[i]) i = id[i];
return i;
}
public boolean connected(int p, int q) {
return root(p) == root(q);
}
public void union(int p, int q) {
int i = root(p);
int j = root(q);
id[i] = j;
}
}