How would the weighted quick-union Algorithm be implemented?

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

Answer

To perform weighted union, you need to know weight of every tree, so make parallel array wt[], where wt[k] contains size of tree with root k. Initial weigths are 1.

Glue smaller tree to the root of larger tree and update weight

public void union(int p, int q) {
    int i = root(p);
    int j = root(q);
    if wt[i] < wt[j] { 
        id[i] = j;
        wt[j] += wt[i] 
    }
    else {similar for j->i}
}

Initialization

public class QuickUnionUF {
  private int[] id;
  private int[] wt;
  public QuickUnionUF(int N) {
    id = new int[N];
    wt = new int[N];
    for (int i = 0; i < N; i++) {
       id[i] = i;
       wt[i] = 1;
    }
  }