Union by rank

A find goes up from an element to its root, so its cost is the height of the tree. Union decides which root goes below the other. Always putting the smaller tree below keeps every tree’s height at most log₂ n.

How trees grow tall

The code of lesson 1 always makes the first root point to the second. Step through the first variant: the unions union(0, 1), union(0, 2), …, union(0, 7) each put the whole tree so far below a single new element. The result is one path, and a find from its bottom goes up through n − 1 parents. Repeat that find n times and the work is about n², no better than the label list.

The mistake is putting a big tree under a small one. Every element of the big tree gets one step further from the root.

The smaller tree goes below

Union by rank gives every root a number called its rank, which starts at 0. When two roots join:

• if their ranks differ, the root with the smaller rank points to the other one, and no rank changes;
• if their ranks are equal, either one can go below, and the rank of the root that stays on top goes up by 1.

Without path compression (lesson 3), the rank of a root is exactly the height of its tree. Joining a shorter tree under a taller one leaves the height as it was. Only two trees of equal height make a taller one, and only by 1. In the figure, the same unions that made a path now make a tree of height 1.

def __init__(self, n):
    self.parent = list(range(n))
    self.rank = [0] * n

def union(self, a, b):
    ra, rb = self.find(a), self.find(b)
    if ra == rb:
        return False
    if self.rank[ra] > self.rank[rb]:
        ra, rb = rb, ra            # ra has the smaller rank
    self.parent[ra] = rb
    if self.rank[ra] == self.rank[rb]:
        self.rank[rb] += 1
    return True

Why the height is at most log₂ n

Here 2^r means r twos multiplied together (and 2⁰ is 1). Proof by induction on the number of unions. At the start every root has rank 0 and one element, and 2⁰ = 1. A union of roots with different ranks does not change any rank, and the root on top only gains elements. A union of two roots that both have rank r makes a root of rank r + 1, and its tree holds both trees: at least 2^r + 2^r = 2^r+1 elements.

There are only n elements, so 2^r ≤ n, which means r ≤ log₂ n. The height of a tree is never more than the rank of its root, so every find goes up at most log₂ n parents: 20 for a million elements. The third variant shows the tallest case for 8 elements: unions of equal pairs, then equal fours, give height 3 = log₂ 8.

Union by size

Instead of a rank, a root can store the size of its tree, the number of its elements. The root of the smaller tree goes below, and the sizes are added. The same proof works: an element moves one level further from its root only when its tree is joined to one at least as large, so its tree at least doubles each time, and that can happen at most log₂ n times.

def union(self, a, b):
    ra, rb = self.find(a), self.find(b)
    if ra == rb:
        return False
    if self.size[ra] > self.size[rb]:
        ra, rb = rb, ra            # ra is the smaller tree
    self.parent[ra] = rb
    self.size[rb] += self.size[ra]
    return True

Size has a bonus: size[find(x)] is the number of elements in x’s set, which many problems ask for. Practice DSUF-P2 needs it. Rank uses less memory in theory (a rank never exceeds log₂ n, so it fits in a byte), which does not matter in Python.