Union-find deeper in Python: union by size and path halving

Why this matters

Union-find answers one question quickly: are these two things in the same group? It also joins groups when you tell it two things belong together. The basic version from the earlier lesson, 05-union-find-and-minimum-spanning-trees, is correct, but its find can walk a long chain of pointers before it gets an answer. Calling find on every element of a chain of one million elements takes about half a trillion hops. This lesson shows how that happens, proves a bound that stops it, and measures the two rules that give the bound.

You need to be able to write a function that uses a list and a while loop, and a class with methods. Every output block in this lesson is the real output of the program shown above it, run with Python 3.14. The random inputs use fixed seeds, so you should see the same numbers on the same Python version. Timings depend on your machine.

The basic version

Union-find stores each element in a list named parent. The entry parent[x] says which element x points to. A root is an element that points to itself. A group is a tree of pointers that ends at one root, and the root names the group.

• find(x) follows the pointers from x until it reaches a root, and returns that root. Each step along a pointer is called a hop.
• union(a, b) finds the root of a and the root of b. If they are the same root, the two elements are already in one group, so nothing changes. If they differ, it makes one root point at the other, so the two groups become one.

The depth of an element is the number of hops from that element to its root. A root has depth 0. The cost of find is the depth of the element you ask about, so depth is the number to keep small.

A chain that makes find slow

Start with the simplest union rule. It makes the root of the first argument point at the root of the second. Take five elements, numbered 0 to 4, and join neighbours in order: union(0, 1), then union(1, 2), then union(2, 3).

Each element starts as its own root, so parent begins as [0, 1, 2, 3]. Here is what happens:

• union(0, 1): the root of 0 is 0, and the root of 1 is 1. Point 0 at 1. Now parent is [1, 1, 2, 3].
• union(1, 2): the root of 1 is 1, and the root of 2 is 2. Point 1 at 2. Now parent is [1, 2, 2, 3].
• union(2, 3): the root of 2 is 2, and the root of 3 is 3. Point 2 at 3. Now parent is [1, 2, 3, 3].

Look at element 0. It points to 1, which points to 2, which points to 3, which is the root. The pointers form a straight line, which computer scientists call a chain. Element 0 is three hops from the root. With a million elements joined in order, element 0 would be a million hops from the root.

The program below builds that chain for 10,000 elements. Then it calls find on every element, from 0 up to 9,999. It counts the hops in four versions of the code: the basic version, the basic version with each of the two new rules, and both rules together. The two rules are called union by size and path halving. The next sections explain each one.

The output is:

no size, no halving    build hops=           0  find hops=    49,995,000
size only              build hops=       9,998  find hops=         9,999
halving only           build hops=           0  find hops=        24,701
size and halving       build hops=       9,998  find hops=         9,999
random unions by size, n=100,000: deepest node depth=8, log2(n)=16

The first row is the chain problem in numbers. The 10,000 finds cost 49,995,000 hops in total. That number is exactly 10,000 times 9,999 divided by 2. Element 0 costs 9,999 hops, element 1 costs 9,998, and so on down to element 9,998, which costs one hop. The total of that countdown is the sum of 1 through 9,999.

The other rows show that both rules remove the problem on this input. The last line of the output is a check on the first rule, and the next section explains it.

The class uses two flags, by_size and halve, and a counter, hops. The counter is only for this lesson. The class at the end of the lesson has no counter.

Optimisation one: union by size

The first rule is simple. When you join two groups, put the root of the smaller group under the root of the larger group. Keep a second list, size, that stores how many elements each group has. The count is only kept correct at roots. A root's entry is the size of its tree. Other entries are not used.

Repeat the same three joins as before, but now use the rule: union(0, 1), then union(1, 2), then union(2, 3).

• union(0, 1): both groups have size 1. On a tie, the code keeps the first root and puts the second under it. So 1 points at 0, and the root 0 now has size 2. parent is [0, 0, 2, 3].
• union(1, 2): the root of 1 is 0, with size 2. The root of 2 is 2, with size 1. The smaller group goes under the larger one, so 2 points at 0, and the size of 0 becomes 3.
• union(2, 3): the root of 2 is 0, with size 3. The root of 3 is 3, with size 1. So 3 points at 0, and the size of 0 becomes 4.

The final parent list is [0, 0, 0, 0]. Every element is one hop from the root. The chain that cost 9,999 hops for the last element is gone.

Why the smaller tree goes under the larger one

The reason is a counting argument. It gives a hard limit on depth, and the limit does not depend on the order of the joins.

Pick one element x. Its depth can only go up in one situation: its root gets attached under another root. Suppose that happens. Before the join, the tree that holds x has some size s. Because the rule attaches the smaller tree, the other tree has at least s elements. After the join, the tree holding x has at least s + s, which is 2s elements.

So every time the depth of x goes up by one, the size of the tree holding x at least doubles. The tree starts with 1 element, when x is alone. It can never hold more than n elements, where n is the total number of elements. If the depth of x is d, the tree holding x has doubled d times, so it has at least 2^d elements. Therefore 2^d is at most n, and d is at most log2 n. The notation log2 n means "the power of two that gives n". For example, log2 8 is 3, because 2 times 2 times 2 is 8.

Depth can also stay the same while a tree grows. Those joins do not hurt the bound, so the bound holds for every sequence of joins.

Here is the bound checked on a real example, using the last line of the first program. With 100,000 elements and 99,999 random joins by size, the deepest element has depth 8. The bound is log2 100,000, which is about 16.6, so the program prints 16 for it. Eight is well under that limit. The exact depth depends on the random joins, so a different seed can give a different depth, but it cannot pass the bound.

The bound has a cost. Every find can still take up to about 16 hops, and it pays that cost again on every call. The next rule removes that repeated cost.

Optimisation two: path halving

The second rule changes find. While it walks up from x, it makes each element point to its grandparent, the element two steps up. Then it moves to that grandparent. The code is one extra line inside the loop:

def find(parent, x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]
        x = parent[x]
    return x

The name comes from the effect. Each pass removes every other element from the path, so the path is roughly halved. Some references call a version that makes every element point straight at the root path compression. Path halving does less work per step and needs only one pass, so it is simpler to write, and it gives the same kind of speed-up.

A worked trace

Start with six elements in a straight chain: 0 points to 1, 1 points to 2, 2 points to 3, 3 points to 4, 4 points to 5, and 5 is the root. In the list, parent is [1, 2, 3, 4, 5, 5]. The root of 0 is 5, and the depth of 0 is 5.

Call find(0):

• The first pass sees x equal to 0. Its parent is 1, which is not the root. The line parent[x] = parent[parent[x]] makes 0 point at the parent of 1, which is 2. Then x moves to 2.
• The second pass sees x equal to 2. Its parent is 3, so the line makes 2 point at 4. Then x moves to 4.
• The third pass sees x equal to 4. Its parent is 5, so the line makes 4 point at the parent of 5, which is 5 itself. Then x moves to 5.
• The loop stops, because 5 points to itself. The root is 5, and the walk took 3 hops.

The walk took 3 hops, where the plain walk takes 5. After the call, 0 points at 2 and is three hops from the root instead of five. Element 2 points at 4, and element 4 points at 5. The next find on any of these elements is cheaper.

The next program runs this trace in code. It returns the root and the number of hops:

chain before: [1, 2, 3, 4, 5, 5]
find(0) returns 5 after 3 hops
chain after: [2, 2, 4, 4, 5, 5]

On a chain of 10,000 elements, the first call still walks about half the chain. Each hop moves two steps up the original chain, and the cheaper walks come only after the first pass has shortened the path. The benefit appears over many calls, which is the topic of the next section.

Measuring both rules together

The first experiment used a chain, which is a worst case that a careless join order can produce. Real inputs are usually less tidy. The next program joins random pairs, then calls find on random elements. It runs four versions. The random joins and the random finds use fixed seeds, so the numbers repeat.

This program uses 20,000 elements and 100,000 finds, which is smaller than the numbers in the Java track. The plain version without halving does over 100 million hops here, and each hop costs more time in Python than in a compiled language. The smaller size keeps the run to a few seconds.

The output is:

naive, no halving    build hops=  9,247,442  find hops=   114,155,800  worst single find=  2,078
naive, halving       build hops=     46,240  find hops=        89,000  worst single find=      4
size, no halving     build hops=     35,120  find hops=       178,823  worst single find=      6
size, halving        build hops=     20,562  find hops=        84,009  worst single find=      2

Divide the find hops by 100,000 to get the average per call. Plain union with no halving averages about 1,142 hops per find, and its worst single call walked 2,078 hops. Union by size alone averages about 1.8 hops, with a worst case of 6. Halving on the naive joins averages about 0.89 hops, but its worst single call was 4, and the proof for it is weaker. Union by size with halving averages about 0.84 hops, and its worst single call was 2.

The build column counts the hops spent inside union, because union calls find twice. Plain union spends about 9.2 million hops just setting up the joins, while the version with both rules spends about 20,000.

On this random input, halving alone does almost as well as both rules together. The chain example shows where that stops being true. Union by size is the part that gives a guarantee for any input, which the next section explains.

Why both rules are needed

Each rule does one job, and the proofs need both jobs.

• Union by size alone keeps every depth at most log2 n, so a single find can never walk more than 16 hops for 100,000 elements. But it does not shorten paths. A deep element costs the same depth every time you ask for it.
• Path halving alone shortens paths every time they are walked, so repeated finds become cheap. But it does not control the shape of the trees. On a chain built by plain joins, the first find still walks about half of the chain. The known proofs of the fast bound need union by size or union by rank as well, and this lesson does not reproduce them.
• Both together give each find a hard limit on the first call, and the work shrinks on later calls. Across many operations, the average cost per operation is close to constant.

The phrase "close to constant" has a precise meaning. Computer scientists write the bound with the inverse Ackermann function, written as the Greek letter alpha with n in brackets. The Ackermann function grows so fast that its inverse grows extremely slowly. For every value of n that a computer can hold, the inverse Ackermann function is at most 4. Robert Tarjan proved a bound of this form for union by rank with path compression, which is close to the two rules here, and later work covers the version with path halving. The full proof is in the research literature and is outside this lesson. The part you can check on paper is the depth bound of log2 n, which the counting argument above proves.

The word "amortised" describes the average. An amortised bound says that the total cost of a sequence of operations is small, even if a single operation in the middle is slow. The first find on a long chain can be slow, and the cost is paid back by the many fast calls that follow.

Writing the class you will reuse

The class below is the version to keep. It has union by size and path halving, and it has no counters, so it is the shape you would use in a real graph program. Compare it with the DSU class in the first program. This version always uses both rules, has no counter or flags, and adds connected and group_size.

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n
 
    # Path halving: every element on the walk points to its grandparent.
    def find(self, x):
        parent = self.parent
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
 
    # Union by size: the smaller root goes under the larger root.
    def union(self, a, b):
        ra = self.find(a)
        rb = self.find(b)
        if ra == rb:
            return False
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        return True
 
    def connected(self, a, b):
        return self.find(a) == self.find(b)
 
    def group_size(self, x):
        return self.size[self.find(x)]

The union method returns True when two groups were joined and False when the two elements already shared a group. Graph algorithms such as Kruskal's minimum spanning tree use that return value to decide whether to keep an edge.

The class stores its data in self.parent and self.size. Each method gets the instance as self, so you call it as uf.find(3) with no list passed in.

Common mistakes

Reading the size of a non-root. The size entry is kept correct only at roots. Writing if self.size[a] < self.size[b]: with a and b as ordinary elements compares stale numbers. The program still runs, but the depth bound no longer holds, and the symptom is slow find calls on large inputs. The fix is to call find first and compare self.size[ra] with self.size[rb], as the class above does.

Forgetting to add the sizes. If the line self.size[ra] += self.size[rb] is missing, the code runs and does not report an error, but the sizes stay at 1. Every comparison is then a tie, so the second root always goes under the first, which can still build a chain if the joins come in the right order. Nothing reports an error, so check the sizes with a small test that prints group_size after a few joins.

Joining elements instead of roots. Writing parent[a] = b instead of making the roots point at each other breaks groups that were already joined. Start with three elements, run parent[0] = 1 and then parent[0] = 2. The list becomes [2, 1, 2]. Element 0 now points at 2, so find(0) returns 2, while find(1) returns 1. Elements 0 and 1 were joined by the first assignment, and the second assignment silently separated them. This program shows it:

parent = [0, 1, 2]
parent[0] = 1
parent[0] = 2
 
 
def find(parent, x):
    while parent[x] != x:
        x = parent[x]
    return x
 
 
print(parent, "find(0) =", find(parent, 0), "find(1) =", find(parent, 1))

[2, 1, 2] find(0) = 2 find(1) = 1

Recursive find on a long chain. A recursive find is shorter to write, but each call takes a little memory on the call stack. Python limits how deep the calls can nest. The limit is 1000 by default, and sys.getrecursionlimit() returns it. A chain of 10,000 elements needs 10,000 nested calls, which is over the limit:

def find_recursive(parent, x):
    if parent[x] == x:
        return x
    return find_recursive(parent, parent[x])
 
 
chain = list(range(1, 10000)) + [9999]
find_recursive(chain, 0)

Running it stops with a RecursionError. The traceback repeats the same line many times before the error, so the end of the message is what matters:

  File "main.py", line 4, in find_recursive
    return find_recursive(parent, parent[x])
  [Previous line repeated 996 more times]
RecursionError: maximum recursion depth exceeded

The file name depends on where you saved the program. Union by size fixes this at the source. With 100,000 elements, the depth is at most log2 of 100,000, which is about 16.6, so a recursive find would never go deeper than 16 calls. Raising the limit with sys.setrecursionlimit is a poor fix, because a chain can always be longer than the limit you set. Keep the loop version anyway, because it uses no extra call frames at all.

Practice

Exercise 1. Five elements start as their own roots. Use union by size, where a tie sends the second root under the first, and a plain find with no halving. Run these joins in order: union(0, 1), union(2, 3), union(0, 2), union(4, 0). What is the final parent list, and how many hops does find(3) take?

Show the answer

After union(0, 1), both sizes are 1, so 1 points at 0. After union(2, 3), 3 points at 2. After union(0, 2), both roots have size 2, so 2 points at 0 and the size of 0 becomes 4. After union(4, 0), the root 4 has size 1 and the root 0 has size 4, so 4 points at 0. The final list is [0, 0, 0, 2, 0]. Element 3 points at 2, and 2 points at 0, so find(3) takes 2 hops.

Exercise 2. Using the counting argument, what is the largest depth possible for a tree of 9 elements built with union by size? Explain your answer without running code.

Show the answer

Depth d needs a tree with at least 2^d elements. A depth of 3 needs at least 8 elements, which 9 allows. A depth of 4 needs at least 16 elements, which 9 does not allow. So the largest depth is 3. The formula gives log2 9, which is about 3.17, and the largest whole number below that is 3.

Exercise 3. Six elements form a chain with parent equal to [1, 2, 3, 4, 5, 5]. Trace one call to find(0) with path halving. Give the final parent list and the number of hops.

Show the answer

The walk visits 0, then 2, then 4, then reaches the root 5. That is 3 hops, not 5. The final list is [2, 2, 4, 4, 5, 5], which is the output of the trace program in the path halving section. Element 4 still points at 5, because the parent of 5 is 5 itself, so the update leaves it unchanged.

Exercise 4. A student writes this union for a graph program: parent[a] = b, with no find calls. Start with three elements, then run union(0, 1) and union(0, 2) with this code. Say what find(1) and find(0) return, and explain why the answer is wrong.

Show the answer

The first call sets parent[0] to 1. The second call sets parent[0] to 2, which overwrites the first join. Now find(0) walks from 0 to 2 and returns 2, while find(1) returns 1. Elements 0 and 1 were joined, so they should return the same root. The code joins elements instead of roots, and the second join overwrote the first.

Exercise 5. A program has 100,000 elements and uses union by size with a recursive find. Explain why this combination is safe, and explain what would go wrong with plain joins and a recursive find.

Show the answer

With union by size, the depth of any element is at most log2 of 100,000, which is about 16.6. A recursive find therefore makes at most 16 nested calls, which is far below Python's limit of 1000. With plain joins on a chain, the depth can reach almost 100,000. A recursive find on that chain makes about 100,000 nested calls, and Python stops the program with a RecursionError, as the common mistakes section shows.

Summary

• Plain union-find can build a chain of pointers, and then one find can walk almost every element.
• Union by size puts the smaller root under the larger root. Each depth increase at least doubles the tree, so the depth stays at most log2 n.
• Path halving makes each element on the walk point to its grandparent. Repeated calls get cheaper, and the first call shortens the path for later calls.
• Both rules together give an average cost per operation that is close to constant, measured with the inverse Ackermann function.
• In the random test, union by size with halving averaged about 0.84 hops per find. Plain union averaged about 1,142.
• A recursive find is unsafe on long chains, and union by size is what keeps the depth small.