Problem Statement

The find function, if called on a child node, may decrease the height of the tree it belongs to, and consequently decrease the rank of the root node, but the algorithm doesn't reflect that update in rank.

Example

As we all know find(0) function will update the root for all nodes from 0 to 3,

	   3               3 
      /              / | \
     2    find(0)   0  1  2
    /      -->     
   1
  / 
 0

rank[3] = 4        rank[3] = 4 still

rank[3] is not updated, it equals to 4 before and after we call find(0), even though the actual height of the tree is 2 afterwards.

Consquence

This could prevent us from finding the most optimized union path.

Example:

Now say if there is another tree with rank[4] = 3

         4          3
        /         / | \
      5          0  1  2
     /
   6

rank[4] = 3    rank[3] = 4

If we call union(4, 3), with the stale rank value rank[3]=4 we will change root[4] to 3, and end up with below merged tree with rank[3] = 4

        3
    / /  |  \
   1  2  0  4
           /
          5
         /
        6
	

However, if we use the correct rank rank[3]=2, we should set the root of 3 to 4, and get the most-optimized output is below with rank[4]=3

      4
    /   \ 
   5     3
  /   /  \  \
 6    0  1   2

Solution

Using size of the tree is a better idea, we don't care about the height of it, but the totally number of nodes under a root.
Because if this root is attacted as a child to another larger root, all its children need to update their root to new larger root when find() is called on them.
Union function is now:

def union(self, A, B):
	root_A = self.find(A)
	root_B = self.find(B)
	if root_A != root_B:
		if self.size[root_A] < self.size[root_B]:
			self.root[root_A] = root_B
			self.size[root_B] += self.size[root_A]
		else:
			self.root[root_B] = root_A
			self.size[root_A] += self.size[root_B]

Notice: The tree in this example may not exist, but for the simplicity of this example, we only care about the height and rank of the tree, but not its full strucutre.