Union-Find (Disjoint Sets)

The problem union-find solves

Imagine a network of computers. Cables get added one at a time, and every so often you need to answer a single question: are these two particular computers connected, whether directly by one cable or indirectly through a chain of other machines? Or imagine a social network where people become friends, and you want to know if two people are in the same friend group. Or a map of roads where towns get linked, and you ask whether you can drive from one town to another.

All of these are the same abstract task. You have a collection of items. Over time some items get joined together into groups. At any moment you want to ask: are these two items in the same group? And when two separate groups get joined, you want to merge them cheaply.

A set here means a group of items with no notion of order and no duplicates. Two sets are disjoint when they share no items at all. The structure in this lesson maintains a family of disjoint sets that together cover every item, and lets groups merge over time. That is why it carries two names: union-find (after its two operations) and disjoint set union, abbreviated DSU.

The naive approach is to store, for each item, a label naming its group, and to merge two groups by walking through every item of one group and relabeling it. Relabeling can touch a large fraction of all the items on a single merge, which is O(n) work per merge where n is the number of items. Union-find does far better: after a small setup, each operation costs so close to constant time that for any input you will ever run, it is effectively a fixed price.

The two operations

Union-find supports exactly two operations, plus the setup.

• find(x): return an identifier for the group that item x currently belongs to. The identifier itself is not meaningful on its own; what matters is that find(a) == find(b) is true exactly when a and b are in the same group.
• union(a, b): merge the group containing a and the group containing b into a single group. If they were already the same group, this does nothing.

With just these two you can answer “are a and b connected?” by checking find(a) == find(b).

The forest representation

The clever idea is to store the groups as a forest. A tree is a set of nodes where every node points to one parent node above it, except one special node at the top that points to itself; that top node is the root. A forest is simply a collection of separate trees. Each tree in our forest is one group, and the root of that tree serves as the group’s identifier.

We do not need real node objects or pointers. Because the items are numbered 0, 1, 2, ..., we can store the whole forest in one ordinary list called parent, where parent[i] is the number of the node that i points to. A list (also called an array) is a numbered sequence of slots. When parent[i] == i, node i points to itself, which is our way of saying i is a root.

At the start every item is alone in its own group, so every item is its own root:

n = 6
parent = list(range(n))   # parent = [0, 1, 2, 3, 4, 5]
# every item is its own parent, so we have 6 separate one-node trees

That leaves six lone trees, each item pointing only at itself:

graph TD
    A0[0]
    A1[1]
    A2[2]
    A3[3]
    A4[4]
    A5[5]

To do find(x), follow the parent links upward until you reach a node that is its own parent. That root is the group identifier. To do union(a, b), find the two roots and make one root point to the other; now both trees share a single root, so they have become one group.

def find(parent, x):
    while parent[x] != x:   # keep walking up until x is a root
        x = parent[x]
    return x

def union(parent, a, b):
    ra = find(parent, a)
    rb = find(parent, b)
    if ra != rb:            # only merge if they are different groups
        parent[ra] = rb     # make root ra point to root rb

The cost of both operations is the cost of find, which walks from a node up to its root. That walk is as long as the tree is tall, so keeping the trees short is what makes the operations fast. The next two ideas do exactly that.

Union by rank or size

If you always attach roots carelessly, you can build a tall, stringy tree: 0 under 1, 1 under 2, 2 under 3, and so on, forming a single chain of height n. A find on the bottom node then walks up all n links, which is O(n). We want the opposite: bushy, shallow trees.

The fix is to be deliberate about which root becomes the child during a union. Two common rules, either works:

• Union by size: attach the root of the smaller tree (fewer nodes) under the root of the larger tree. The larger tree’s height does not grow; the smaller one only gains one level.
• Union by rank: rank is an estimate of a tree’s height. Attach the lower-rank root under the higher-rank root. Only when the two ranks are equal does the surviving root’s rank increase by one.

Both rules keep every tree’s height at O(log n) even without the next optimization, because a tree of height h built this way must contain at least 2**h nodes, so height cannot exceed log2 of the item count.

Watch union by size on a small example. Beforehand there is a two-node group {3, 4} rooted at 4, sitting next to the lone node 2:

graph TD
    B4[4] --> B3[3]
    B2[2]

Because {3, 4} is the larger tree, union(2, 3) hangs 2 under its root 4 rather than pulling 4 under 2:

graph TD
    C4[4] --> C3[3]
    C4 --> C2[2]

Node 2 is now one step from the root, and the tree stayed short.

Path compression

The second optimization works during find itself. When you walk from a node up to its root, you learn the root of every node you stepped through along the way. Path compression means: after finding the root, point each of those nodes directly at the root, so next time the walk is a single step.

Before find(0), suppose the group is a chain: 0 points to 1, 1 points to 2, 2 is the root.

graph TD
    D0[0] --> D1[1]
    D1 --> D2[2]

After find(0) with path compression, both 0 and 1 point straight at the root 2. The chain has been flattened:

graph TD
    E0[0] --> E2[2]
    E1[1] --> E2[2]

The tree is now as shallow as possible, so every future find on these nodes is nearly instant. Here is find with path compression written cleanly with recursion. Recursion means a function that calls itself on a smaller piece of the same problem.

def find(parent, x):
    if parent[x] != x:              # x is not the root
        parent[x] = find(parent, parent[x])   # set x's parent to the root
    return parent[x]

The line parent[x] = find(...) is where compression happens: on the way back down the recursion, every node visited gets its parent rewritten to the root.

A complete DSU class

Here is a full, runnable disjoint-set structure using both union by size and path compression. A class is a blueprint that bundles data (here the parent and size lists) with the operations that act on it. self refers to the particular structure you are working with.

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))   # each item starts as its own root
        self.size = [1] * n            # every group starts with 1 node

    def find(self, x):
        # path compression: flatten the walk to the root
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False               # already together; nothing merged
        # union by size: hang the smaller tree under the larger root
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra            # make ra the larger root
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        return True                    # a real merge happened

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


dsu = DSU(6)
dsu.union(0, 1)
dsu.union(2, 3)
dsu.union(1, 3)                        # merges the group {0,1} with {2,3}
print(dsu.connected(0, 2))             # -> True
print(dsu.connected(0, 4))             # -> False
print(dsu.union(0, 3))                 # -> False   (already connected)

union returns True when it actually merged two different groups and False when the two items were already together. That return value is directly useful, as the cycle-detection use case below shows.

Step-by-step trace

Start with 6 items, each alone: groups are {0} {1} {2} {3} {4} {5}. We run a sequence of unions and watch the parent list and the groups change. This trace uses union by size; ties in size attach the second root under the first.

Read step 3 carefully, since it is the interesting one. union(1, 3) first calls find(1), which walks 1 -> 0 and returns root 0, then find(3), which walks 3 -> 2 and returns root 2. The two groups have equal size (2 each), so root 2 is attached under root 0, giving parent[2] = 0. Notice that node 3 still points at 2, not directly at 0; its link is only shortened later, the first time a find passes through it and path compression rewrites it.

Use cases

Connected components. Given a set of items and a list of pairwise links, how many separate groups are there, and which items share a group? Run union on every link, then count the distinct roots. This is the canonical use.

edges = [(0, 1), (2, 3), (1, 3), (4, 5)]
dsu = DSU(6)
for a, b in edges:
    dsu.union(a, b)
roots = {dsu.find(i) for i in range(6)}   # a set of the distinct group roots
print(len(roots))                          # -> 2   (the groups {0,1,2,3} and {4,5})

Cycle detection in an undirected graph. A graph is a set of nodes joined by edges (links). A cycle is a loop: a path that returns to where it started. Process the edges one at a time; for edge (a, b), if a and b are already in the same group, then adding this edge closes a loop, so a cycle exists. If they are in different groups, union them and move on. Because our union returns False when the items were already connected, the check is one line.

def has_cycle(n, edges):
    dsu = DSU(n)
    for a, b in edges:
        if not dsu.union(a, b):   # union failed => a and b were already linked
            return True
    return False

print(has_cycle(3, [(0, 1), (1, 2), (0, 2)]))   # -> True   (the triangle is a loop)
print(has_cycle(3, [(0, 1), (1, 2)]))           # -> False  (a plain chain, no loop)

Kruskal’s algorithm for the minimum spanning tree. A spanning tree connects all nodes of a graph using a subset of edges and no cycles; the minimum one has the smallest total edge weight. Kruskal sorts the edges from cheapest to most expensive and adds each edge only if it joins two currently separate groups (checked with union-find). Skipping any edge whose endpoints are already connected is exactly the cycle check above, and it is what keeps the result a tree.

Big-O table

Let n be the number of items and m the number of operations. α(n) is the inverse Ackermann function, explained just below.

Why the fast bound holds, in plain words: union by size or rank guarantees the trees never grow taller than log n, and path compression flattens every path it touches so repeated finds get cheaper and cheaper. Combined, a proven result says a sequence of m operations costs O(m · α(n)) in total. Amortized means averaged over the whole sequence: one individual find might occasionally walk a few links, but the total across all operations is bounded, so the average per operation is O(α(n)).

The inverse Ackermann function α(n) grows so slowly that it is at most 4 for any number of items that could physically exist (far beyond the number of atoms in the observable universe). So although it is not literally a constant, in every practical sense each operation is constant time. That is the headline result of this structure: near-constant amortized time per operation.

Common pitfalls

• Comparing roots, not the raw items. find(a) == find(b) answers “same group”; a == b only answers “same item”. Always compare the results of find.
• Forgetting to check equal roots in union. If you skip the if ra == rb guard and then execute parent[ra] = rb when ra == rb, you set a root to point at itself, which is harmless here, but with union by rank you can wrongly bump a rank. Always return early when the roots match.
• Union by height without path compression. Path compression changes actual heights, which makes true height hard to track. Union by rank deliberately keeps rank as an upper-bound estimate and never decreases it, so it stays correct alongside compression. Do not try to maintain exact height.
• Recursion depth on huge inputs. The recursive find is clean, but before path compression has run, a pathological chain could be deep enough to exceed Python’s recursion limit. Union by size/rank prevents deep chains from forming in the first place; if you are nervous, write find as a loop with a second pass that rewrites parents.
• Reusing one DSU across independent problems. Each problem needs a fresh DSU(n). Carrying old unions into a new problem silently reports items as connected when they are not.

Practice

1. Add a method count(self) to the DSU class that returns how many separate groups currently exist. Do it in two ways: once by counting distinct roots with find, and once by maintaining a running counter that starts at n and decreases by one on every successful union.

2. Given n = 7 and the edge list [(0,1), (1,2), (3,4), (5,6), (2,4)], work out by hand the final groups and the number of connected components, then write code to confirm your answer.

3. Rewrite find as a loop instead of recursion, still performing path compression. Hint: walk up once to locate the root, then walk up a second time rewriting every node’s parent to that root.