Disjoint sets

Some problems split things into groups that only ever join together, never split apart: friends of friends, connected places on a map, pixels of one colour. Union-find keeps track of such groups with two operations, and both are almost free.

Groups that only merge

Two sets are disjoint when they have no element in common. A collection of disjoint sets is also called a partition: every element is in exactly one set. Number the elements 0 to n − 1. At the start each element is in a set of its own. We want two operations:

• find(x): which set is x in? The answer is one chosen element of the set, called its representative. The name itself does not matter; what matters is that find(x) == find(y) exactly when x and y are in the same set.
• union(a, b): join the set of a and the set of b into one.

The structure is called union-find, or a disjoint-set union (DSU). It cannot split a set again, and it cannot list a set’s elements quickly. It does one thing, and does it very fast.

A label per element

The first idea is a list label where label[x] is the representative of x’s set. Then find is one read. But union must change the label of every element of one set, and the only way to find them all is to read the whole list. Step through the second variant of the figure: every union costs n reads, so n unions cost n² (n times n).

A forest of parent pointers

The better idea is to store each set as a tree, where every element points to a parent in the same set, and the root points to itself. The root is the representative. A collection of trees is called a forest. One list stores the whole forest: parent[x] is the parent of x.

• find(x): follow parents up from x until an element is its own parent. Return it.
• union(a, b): find the two roots. If they are different, make one root point to the other. That is a single change to the list.

Why does this stay true? A union changes only the parent of a root. Every element that reached the old root now reaches the new root, one step further up, and nothing else changes. So the two sets become one, and every other set is untouched.

Unlike the trees of TREE and BALT, the order inside these trees means nothing, and a node may have any number of children. Only the root matters. The links also point upwards, from child to parent, so a node never needs to know its children.

The code

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))   # everyone is a root

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

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False               # already in one set
        self.parent[ra] = rb
        return True

    def connected(self, a, b):
        return self.find(a) == self.find(b)

Returning False when the two elements are already together is a small, useful habit: lesson 5 uses it to find cycles. list(range(n)) builds [0, 1, …, n − 1], so every element starts as its own parent.

A union now costs two finds and one change. But a find costs the height of the tree, and nothing so far stops a tree from growing into one long path of n nodes. The next two lessons add two simple rules that keep the trees almost flat.