Union-Find — Two Optimizations Make It Practically Constant

Graph problems often ask the same question repeatedly: are these two elements in the same group? How many groups remain when friends of friends are linked together, and does adding one more edge create a cycle. While organizing the graph topic I implemented the standard tool for this question, Union-Find (Disjoint Set), starting from the version with no optimizations, and got to see why the two optimizations always come attached. Without them, find stretches to O(N); with both, it becomes practically constant.

find and union

Give each group one representative (root), and make every element reach its group’s root by following parent pointers. Two elements are in the same group when their roots match. find(x) follows parents from x and returns the root; union(x, y) finds both roots and attaches one to the other.

class UnionFindNaive:
    def __init__(self, n):
        self.parent = list(range(n))

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

    def union(self, x, y):
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False
        self.parent[root_y] = root_x
        return True

This implementation produces correct answers. The problem is the shape of the tree. Calling union(1, 0), union(2, 1), union(3, 2), union(4, 3) in that order makes each new element adopt the existing tree’s root as its child, producing a single chain: 0 → 1 → 2 → 3 → 4. Now find(0) climbs four steps to the root. With N elements, find degrades to O(N) in the worst case, and since union calls find internally, N repetitions cost O(N²).

flowchart BT
    c0["0"] --> c1["1"] --> c2["2"] --> c3["3"] --> c4["4"]

Path Compression

The first optimization lives inside find. While climbing to the root, rewrite the parent of every node on the path to point directly at the root. From the next find on, all of those nodes reach the root in a single step.

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

The recursion climbs to the root and, on the way back, attaches every node on the path directly to it. One call to find(0) on the chain above reshapes the tree like this.

flowchart BT
    f0["0"] --> f4["4"]
    f1["1"] --> f4
    f2["2"] --> f4
    f3["3"] --> f4

The more a tree is queried, the faster it flattens.

Union by Rank

The optimization on the union side decides the direction in which trees are attached. Attaching arbitrarily can produce chains, so attach the shorter tree to the taller one. The overall height grows by one only when the two trees are equally tall.

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n

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

    def union(self, x, y):
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False
        if self.rank[root_x] < self.rank[root_y]:
            root_x, root_y = root_y, root_x
        self.parent[root_y] = root_x
        if self.rank[root_x] == self.rank[root_y]:
            self.rank[root_x] += 1
        return True

The rank is an upper bound on the tree height before compression. Path compression shortens the real height without updating the rank, and that is still sufficient as a criterion for choosing the merge direction. Using group size instead of rank — attaching the smaller tree to the larger — gives the same guarantee, and the size variant wins when group sizes need to be queried.

With both optimizations, find and union run in O(α(N)) amortized. α is the inverse Ackermann function, which stays at 4 or below for every realistic N, so it is treated as practically constant.

Counting Connected Components

Number of Provinces (LeetCode 547) gives friendships as an adjacency matrix and asks for the number of groups. Start the connected-component count at N, decrement whenever a union succeeds, and the value at the end of the scan is the answer.

def count_provinces(is_connected):
    n = len(is_connected)
    uf = UnionFind(n)
    count = n
    for i in range(n):
        for j in range(i + 1, n):
            if is_connected[i][j] and uf.union(i, j):
                count -= 1
    return count

The matrix is symmetric, so scanning the upper triangle (i < j) is enough. union returns False for elements already in the same group, so the count never decreases incorrectly.

That return value extends to cycle detection. Adding edges one by one, the moment union returns False, the two endpoints were already connected, so that edge closes a cycle. Redundant Connection (LeetCode 684) is exactly this pattern.

Selection Criteria

BFS/DFS also handles connected components.

When BFS/DFS fits:

• The graph is fixed and one traversal answers the question
• The contents of a component (paths, visit order) are needed
• Neighbors are computed from coordinates, as in grid traversal

When Union-Find fits:

• Edges arrive incrementally with connectivity queries interleaved
• Each edge addition must be checked immediately for creating a cycle
• An edge must be repeatedly judged for adoption, as in Kruskal’s MST

Whether two elements share a group is settled by a single root representative, and keeping that verdict fast is the job of the two optimizations. Only after implementing the unoptimized version first did I understand why these two always come as a pair.