Union-Find (Disjoint Sets)

Union-find (disjoint set union, DSU) maintains a family of disjoint groups over numbered items, letting you merge groups and ask “are these two in the same group?” in near-constant amortized time.

Core operations

• find(x): return the group’s identifier (its root). find(a) == find(b) iff a and b share a group.
• union(a, b): merge the groups of a and b; does nothing if already together.
• connected(a, b): shorthand for find(a) == find(b).
• Naive relabeling merge is O(n) per union; DSU is far cheaper.

Forest representation

• Store groups as a forest of trees; each tree’s root is the group id.
• One parent list: parent[i] is i’s parent; parent[i] == i marks a root.
• Init: parent = list(range(n)) — every item its own root, n singletons.
• find: follow parent links up to the root. union: point one root at the other.
• Both cost the height of the tree, so keep trees short.

The two optimizations

• Union by size: hang the smaller tree’s root under the larger’s; larger height unchanged.
• Union by rank: rank estimates height; attach lower rank under higher; on a tie the survivor’s rank goes up by one.
• Either rule alone keeps height O(log n) (a tree of height h has ≥ 2**h nodes).
• Path compression (during find): after reaching the root, repoint every node on the walk directly at the root, flattening the chain.

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])  # path compression
        return self.parent[x]
    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra           # ra is the larger root
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        return True

• union returns True on a real merge, False when already connected.

Complexity

• α(n) is the inverse Ackermann function: ≤ 4 for any real input, so effectively constant.
• Amortized: total over m ops is O(m · α(n)); a single find may walk a few links.

Use cases

• Connected components: union every edge, then count distinct roots ({find(i) for i in range(n)}).
• Cycle detection (undirected): for edge (a,b), if not union(a,b) a cycle exists.
• Kruskal’s MST: sort edges cheapest-first, add each edge only if union succeeds (joins separate groups).

Gotchas

• Compare find(a) == find(b), not the raw items a == b.
• Keep the if ra == rb: return guard, or union by rank can wrongly bump a rank.
• Don’t track exact height with path compression; use union by rank (an upper-bound estimate that never decreases).
• Recursive find can hit Python’s recursion limit on a deep chain before compression runs; union by size/rank prevents deep chains, or write find as a loop.
• Use a fresh DSU(n) per problem; reusing one leaks old unions.