Concepts: Disjoint Sets and Union-Find

Union-find is a data structure that keeps a collection of elements split into non-overlapping groups. It is also called the disjoint-set data structure, or DSU (disjoint set union). This chapter covers the problem it solves, its two core operations, and the vocabulary you will meet when studying it.

What is a disjoint set?

Two sets are disjoint when they share no element. Splitting all elements into disjoint sets so that every element belongs to exactly one of them is called a partition. For example, the elements 0 to 5 can be partitioned into {0, 1, 2}, {3, 4} and {5}.

Union-find stores such a partition and is built to answer two questions quickly:

• find(x): which set does x belong to? Concretely, it returns the set's representative element.
• union(a, b): merge the set containing a with the set containing b.

To check whether two elements are in the same set, compare find(a) == find(b). The key trick is that sets have no separate names; one member of each set acts as its representative.

Intuition: merging friend groups

Imagine that everyone at a school starts out alone. Each time news arrives such as "Alex and Beth became friends" or "Beth and Chris became friends", you merge the groups the two people belong to. Alex and Chris end up in the same group even though they have never met. Relationships that chain transitively, connections that only ever get added, and a single question, "same group?", are exactly the setting where union-find shines.

The simplest implementation writes a group label next to every element. find is instant, but every union has to scan the whole array.

# Naive approach (quick-find): one union costs O(n)
label = list(range(6))

def find(x):
    return label[x]

def union(a, b):
    la, lb = label[a], label[b]
    for i in range(len(label)):
        if label[i] == lb:
            label[i] = la

union(0, 1)
union(1, 2)
print(find(2) == find(0))  # True

Representing sets as a forest

Union-find represents each set as a tree and the whole partition as a forest of trees. Every element remembers only its parent; an element that is its own parent is a root, which is the representative of its set. find walks up the parent links to the root, and union simply hangs one root under the other.

parent = list(range(6))  # initially every element is its own root

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

def union(a, b):
    ra, rb = find(a), find(b)
    if ra != rb:
        parent[rb] = ra  # always link roots, never arbitrary nodes

union(0, 1); union(1, 2); union(3, 4)
print([find(x) for x in range(6)])  # [0, 0, 0, 3, 3, 5]

After running the code above, the parent array is [0, 0, 0, 3, 3, 5]. Three slots point to themselves, 0, 3 and 5, so there are three sets. The number of roots is always the number of sets, and every successful merge removes exactly one root.