Union-Find: The One Algorithm to Rule Them All

The Quest Begins (The "Why")

I still remember the first time I stared at a LeetCode problem that asked me to count the number of islands in a grid. My initial instinct was to run a DFS/BFS from every land cell, marking visited cells as I went. It worked… until the test cases grew to 10⁵ × 10⁵ grids and my solution timed out. I felt like I was swinging a wooden sword at a dragon—lots of effort, barely any scratch.

That frustration pushed me to look for a smarter way to answer “are these two nodes in the same component?” without traversing the whole graph each time. The answer came in the form of a modest‑looking data structure called Union‑Find (also known as Disjoint Set Union, DSU). Once I grasped it, the dragon fell in a single swing. Let’s see why it works so well.

The Revelation (The Insight)

At its heart, Union‑Find maintains a collection of disjoint sets and supports two operations:

1. Find(x) – returns a representative (the “root”) of the set containing x.
2. Union(x, y) – merges the sets containing x and y.

If we can quickly tell whether two elements share the same root, we can answer connectivity questions in near‑constant time. The magic lies in two simple heuristics:

• Union by size (or rank) – always attach the smaller tree under the larger one. This keeps the overall depth shallow.
• Path compression – during Find, we make every node on the path point directly to the root. Future finds become exponentially faster.

Why does this give us almost O(1) per operation? Imagine each union as adding an edge to a forest. Without heuristics, a chain could grow to length n, making finds O(n). By always attaching the shorter tree to the taller one, the height of any tree grows at most logarithmically. Path compression then repeatedly halves the remaining distance to the root, leading to the inverse Ackermann function α(n)—which grows so slowly that for any practical n it’s ≤ 5. In interview land we treat it as O(1).

Putting it together, a sequence of m union/find operations on n elements runs in O(n + m α(n)) time, essentially linear. The space usage is just two integer arrays of size n: parent and size.

Wielding the Power (Code & Examples)

Before: Naïve DFS per query

def count_islands(grid):
    rows, cols = len(grid), len(grid[0])
    visited = [[False]*cols for _ in range(rows)]

    def dfs(r, c):
        if r < 0 or r >= rows or c < 0 or c >= cols:
            return
        if grid[r][c] == '0' or visited[r][c]:
            return
        visited[r][c] = True
        dfs(r+1, c); dfs(r-1, c); dfs(r, c+1); dfs(r, c-1)

    islands = 0
    for r in range(rows):
        for c in range(cols):
            if grid[r][c] == '1' and not visited[r][c]:
                islands += 1
                dfs(r, c)
    return islands

Each dfs could walk the whole grid → O(rows × cols) per new island, which blows up on large inputs.

After: Union‑Find solution

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size   = [1] * n          # for union by size

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

    def union(self, x, y):
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return False                # already in same set
        # union by size: attach smaller tree under larger
        if self.size[rx] < self.size[ry]:
            rx, ry = ry, rx
        self.parent[ry] = rx
        self.size[rx] += self.size[ry]
        return True

def num_islands(grid):
    if not grid: return 0
    rows, cols = len(grid), len(grid[0])
    uf = UnionFind(rows * cols)
    # count of land cells initially
    land = sum(cell == '1' for row in grid for cell in row)

    for r in range(rows):
        for c in range(cols):
            if grid[r][c] != '1':
                continue
            idx = r * cols + c
            # only check right and down to avoid double work
            if r + 1 < rows and grid[r+1][c] == '1':
                if uf.union(idx, (r+1)*cols + c):
                    land -= 1
            if c + 1 < cols and grid[r][c+1] == '1':
                if uf.union(idx, r*cols + (c+1)):
                    land -= 1
    return land

What changed?

• We treat each cell as a node in a DSU.
• Every time we union two adjacent lands, we decrement the island counter because two separate components just merged.
• The find operation is amortized almost constant thanks to path compression, and union by size keeps the trees flat.

Common traps to avoid

• Forgetting to apply path compression (leads to tall trees and slower finds).
• Using union by rank incorrectly—always attach the smaller tree under the larger one; swapping the condition creates degenerate cases.
• Not skipping already‑unioned pairs (the if uf.union(...): land -= 1 guard) – otherwise you’d under‑count islands.

Second example: Redundant Connection (LeetCode 684)

The problem gives an undirected graph that started as a tree (n nodes, n‑1 edges) and then one extra edge was added, creating a single cycle. We need to return that edge.

def find_redundant_connection(edges):
    n = len(edges)
    uf = UnionFind(n + 1)      # nodes are 1-indexed
    for u, v in edges:
        if not uf.union(u, v):  # union failed → already connected
            return [u, v]

Here the same Union‑Find class does the heavy lifting: when we try to union two vertices that already share a root, we’ve found the redundant edge. No DFS, no adjacency lists—just a couple of array accesses per edge.

Why This New Power Matters

With Union‑Find in your toolbox, you can:

• Solve dynamic connectivity problems (e.g., “are these two users in the same network?”) in near‑real time.
• Tackle grid‑based challenges like number of islands, surrounded regions, or swapping islands with ease.
• Handle offline queries (process all unions first, answer connectivity later) with the same structure.

The best part? The code is tiny—under 30 lines—but the impact is massive. I’ve seen interviewers light up when a candidate replaces a DFS flood‑fill with a clean DSU solution; it signals that you understand both the theory and the practical tricks that make algorithms fly.

Your Turn

Grab a problem you’ve previously solved with BFS/DFS (maybe “Friend Circles” or “Accounts Merge”) and rewrite it using Union‑Find. Notice how the runtime drops and how the code becomes easier to reason about. Then, try extending it: support deletions with a “rollback” version, or handle weighted unions.

Go forth, wield the One Algorithm, and may your runs be swift and your bugs few! 🚀