Challenge: union-find

Think of a jigsaw puzzle on a table. Pieces snap together one pair at a time, and the loose pieces and finished chunks on the table are the connected components of a graph that only ever grows. In between snaps you may want to know whether two pieces already belong to the same chunk, or how many chunks are left. The DFS of step 3 can tell you, but it starts from nothing every time: O(V + E) for each question, again and again as the puzzle grows. Union-find, also called a disjoint set union (DSU), is built for groups that only merge. It answers "same group?" and joins two groups, each in almost O(1).

A parent list

Each group is a small tree, kept in a list parent: parent[x] is the vertex above x. The vertex at the top of a group, its root, is its own parent, and it names the group. At the start, every vertex is alone, so parent = list(range(n)). Two operations do the work:

• find(x) follows the parents up from x to the root. Two vertices are in the same group exactly when they have the same root.
• union(a, b) finds both roots. If they differ, it hangs one root under the other, and the two groups are one.

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


def union(parent, a, b):
    root_a, root_b = find(parent, a), find(parent, b)
    if root_a != root_b:
        parent[root_a] = root_b


parent = list(range(7))
union(parent, 5, 2)
union(parent, 0, 6)
union(parent, 5, 0)
print(parent)
print(find(parent, 5) == find(parent, 6), find(parent, 1) == find(parent, 3))

[6, 1, 6, 3, 4, 2, 6]
True False

After the three unions, 0, 2, 5 and 6 share the root 6, while 1, 3 and 4 are still alone. The tree needn't look like the edges you added: the last union named 5 and 0, but it hung 5's root, 2, under 0's root, 6, so parent[5] is still 2 and 5 is two steps below the root.

The trap: tall trees

Nothing here keeps the trees short. When person after person joins the group of person 0 with union(0, i), each union hangs the old root under i, and 0 sinks to the bottom of a chain:

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


def depth(parent, x):
    steps = 0
    while parent[x] != x:
        x = parent[x]
        steps += 1
    return steps


for n in [10, 100, 1000]:
    parent = list(range(n))
    for i in range(1, n):
        root_a, root_b = find(parent, 0), find(parent, i)
        if root_a != root_b:
            parent[root_a] = root_b
    print(n, "people: 0 is", depth(parent, 0), "steps below its root")

10 people: 0 is 9 steps below its root
100 people: 0 is 99 steps below its root
1000 people: 0 is 999 steps below its root

Each find(0) walks the whole chain so far, so the n unions take 1 + 2 + ... + (n - 1) steps: O(n²), 5 billion steps for 100,000 people.

Two fixes

• Union by size: keep the size of each group at its root, and hang the root of the smaller group under the root of the bigger one. A vertex gets one step further from its root only when its group at least doubles in size, which can happen at most log₂ n times. So no tree is ever taller than log₂ n.
• Path compression: once find has found the root, it walks the path again and points every vertex on it straight at the root. The next find from any of them takes one step.

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


def union(parent, size, a, b):
    root_a, root_b = find(parent, a), find(parent, b)
    if root_a == root_b:
        return False
    if size[root_a] < size[root_b]:
        root_a, root_b = root_b, root_a
    parent[root_b] = root_a
    size[root_a] += size[root_b]
    return True


parent = list(range(6))
size = [1] * 6
groups = 6
for a, b in [(0, 1), (3, 4), (1, 2), (2, 0), (1, 4)]:
    if union(parent, size, a, b):
        groups -= 1
    else:
        print("redundant:", a, b)
print(groups, "groups")
print(size[find(parent, 3)], find(parent, 0) == find(parent, 4))

parent = list(range(1000))
size = [1] * 1000
for i in range(1, 1000):
    union(parent, size, 0, i)
print(max(parent))

redundant: 2 0
2 groups
5 True
0

size[x] only counts while x is a root, so the size of a group is size[find(parent, x)]. union returns True only when it joins two groups, so the number of groups starts at n and goes down by one with every True. In the second part, where 999 people join person 0's group one at a time, the big group's root 0 stays on top, and every other person hangs right under it: the largest parent is 0.

With both fixes, m operations on n vertices take O(m·α(n)) time, where α is a function that grows so slowly that it's at most 4 for any n that fits in a computer. In practice that's O(1) for each operation, with O(n) memory. Write find with loops, as here: a recursive find is shorter, but on a tall tree it hits Python's recursion limit before it has compressed anything.

Redundant edges

union returns False when a and b are already in the same group. Then the edge between them connects nothing new: it is redundant, and in graph words it closes a cycle. That's how union-find finds the extra cable in a network, and how Kruskal's algorithm builds the cheapest network that connects every vertex: it goes through the edges from the cheapest up and skips each one that union says is redundant.