Union-Find (Disjoint Set)

The question this structure answers efficiently

Many problems reduce to a repeated pattern: "are these two elements in the same group?" and "merge these two groups into one." DFS/BFS (previous lesson) can answer "same group" by re-running a full traversal each time, but that's wasteful when groups merge and get queried over and over — you'd re-traverse from scratch on every query. Union-Find (also called Disjoint Set Union, DSU) is a structure purpose-built for exactly these two operations, find (which group is this element in?) and union (merge two groups), each running in very-nearly O(1) time after two specific optimizations are applied together.

The basic idea: a forest of parent pointers

Picture groups of children on a playground, each group playing a game where everyone holds hands in a chain leading back to whoever is currently "it" — the group's leader. To find out who's leading your group, you don't need a roster; you just follow the hand-holding chain until you reach someone holding nobody's hand but their own. Two children are in the same game exactly when following their chains leads to the same leader.

Represent each group as a tree, where every element points to a parent, and a group's identity is the root of its tree (a root is its own parent). Initially, every element is its own group — every element is its own root:

class UnionFind:
    def __init__(self, n: int) -> None:
        self.parent = list(range(n))    # every element starts as its own root

    def find(self, x: int) -> int:
        while self.parent[x] != x:      # walk up until reaching a root
            x = self.parent[x]
        return x

    def union(self, a: int, b: int) -> None:
        root_a, root_b = self.find(a), self.find(b)
        if root_a != root_b:
            self.parent[root_a] = root_b   # attach one tree under the other's root

find(x) walks parent pointers up to the root; two elements are in the same group exactly when find returns the same root for both. union finds both roots and attaches one under the other. As written, this is correct but can degrade badly: repeatedly unioning in a bad order (e.g. always attaching the new tree under the OLDER one in a long chain) can build a tree that's effectively a linked list, making find cost O(n) in the worst case.

Optimization 1: union by rank/size — keep trees shallow

Instead of arbitrarily attaching one root under the other, track each root's rank (an upper bound on its subtree's height) or size (element count), and always attach the SMALLER tree under the root of the LARGER one. This keeps the resulting tree's height from growing unnecessarily: attaching a smaller tree under a bigger one can increase height by at most 1, and only when the two trees were the exact same size — attaching a bigger tree under a smaller one, by contrast, can add the entire smaller tree's height on top. Using size (simpler to reason about) as the criterion:

class UnionFind:
    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.size = [1] * n              # each root's tree size

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

    def union(self, a: int, b: int) -> None:
        root_a, root_b = self.find(a), self.find(b)
        if root_a == root_b:
            return
        if self.size[root_a] < self.size[root_b]:      # attach SMALLER under LARGER
            root_a, root_b = root_b, root_a
        self.parent[root_b] = root_a
        self.size[root_a] += self.size[root_b]

With union by size/rank alone (no path compression yet), tree height is bounded by O(log n) — a standard result: a tree of height h built this way must have at least 2^h elements (each merge that increases height requires combining two equally-tall trees, doubling the count). Turn that around: with n elements total, 2^h ≤ n, so taking log₂ of both sides gives h ≤ log₂ n — height can be at most log₂ n. This alone already brings find down to O(log n).

Optimization 2: path compression — flatten on the way up

Every time find(x) walks up to the root, it has just learned the true root of every node it passed through along the way. Path compression exploits this: after finding the root, re-point every node visited during that walk DIRECTLY to the root, so the next find on any of those nodes is O(1):

def find(self, x: int) -> int:
    root = x
    while self.parent[root] != root:        # first pass: locate the root
        root = self.parent[root]
    while self.parent[x] != root:           # second pass: flatten the path
        self.parent[x], x = root, self.parent[x]
    return root

(An equivalent, more common recursive form: if parent[x] != x: parent[x] = find(parent[x]). Both achieve the same flattening — every node on the path from x to the root ends up pointing directly at the root.)

Why both together give near-constant time

Union by rank/size alone gives O(log n). Path compression alone also gives an amortized bound close to O(log n) (a more involved argument). Together, the two optimizations combine to give an amortized time per operation of O(α(n)), where α is the inverse Ackermann function — a function that grows so slowly that for any n that could ever exist in practice (far, far beyond the number of atoms in the observable universe), α(n) is at most 4 or 5. This is, for every practical purpose, constant time. The formal proof of this bound is involved and not reproduced here; what matters operationally is the combination — using only one of the two optimizations still leaves a provably worse (though still good, O(log n)) bound, while using neither lets find degrade toward O(n) on adversarial union orders, exactly the linked-list-shaped worst case from the very first, unoptimized version above.

Where this is used

Union-Find is the natural structure whenever a problem is fundamentally about connectivity — do these elements end up grouped together — rather than about paths or distances (which is DFS/BFS/Dijkstra's territory instead). This module's problems apply it to detecting a cycle-causing edge (Redundant Connection), counting connected components (Number of Provinces), and, in the next lesson, building a Minimum Spanning Tree (Kruskal's algorithm uses Union-Find to detect when adding an edge would create a cycle).