The Fellowship of Sets: Union‑Find (Disjoint Set) from Theory to LeetCode

The Quest Begins (The “Why”)

I still remember the first time I stared at a LeetCode problem that asked me to count how many separate groups of friends existed in a social network. My initial instinct was to throw a nested loop at it, compare every pair, and mark visited nodes. The code worked on the tiny examples, but as soon as the input grew to 10⁵ users it felt like I was trying to bail out a sinking ship with a teaspoon—slow, painful, and doomed to timeout.

That frustration sparked a question: Is there a way to merge groups together almost instantly, and then ask “are these two people in the same group?” in constant time? The answer turned out to be a humble data structure called Union‑Find (also known as Disjoint Set Union, DSU). It’s not flashy, but once you grasp why it works, it feels like unlocking a secret shortcut in a game you’ve been grinding for hours.

The Revelation (The Insight)

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

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

The magic lies in how we store the parent links. Imagine each element pointing to its parent, forming a forest of trees. The root of a tree is the set’s identifier. If we always attach the smaller tree under the larger one (union by size/rank), the height of any tree stays logarithmic. Then we add path compression: during a Find, we make every node on the path point directly to the root.

Why does this give us near‑constant time?

• Union by size guarantees that a node’s depth can increase at most O(log n) before it becomes a child of a larger tree.
• Path compression flattens the tree each time we query, so future finds on those nodes become O(1).

The combination yields an amortized complexity of O(α(n)), where α is the inverse Ackermann function—so slow‑growing that for any practical input it’s essentially constant. In plain English: after a few unions and finds, the structure almost “remembers” the answer for you.

Wielding the Power (Code & Examples)

Before: The Naïve Approach (for contrast)

def count_islands_naive(grid):
    visited = set()
    def dfs(r, c):
        if (r, c) in visited or not (0 <= r < len(grid) and 0 <= c < len(grid[0])) \
           or grid[r][c] == '0':
            return
        visited.add((r, c))
        for dr, dc in [(1,0),(-1,0),(0,1),(0,-1)]:
            dfs(r+dr, c+dc)

    islands = 0
    for i in range(len(grid)):
        for j in range(len(grid[0])):
            if grid[i][j] == '1' and (i, j) not in visited:
                dfs(i, j)
                islands += 1
    return islands

This works, but each DFS can revisit cells many times, leading to O(m·n·α) in worst‑case patterns and a lot of recursion depth.

After: Union‑Find Solution (Number of Islands – LeetCode 200)

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

    def find(self, x):
        # path compression
        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
        # 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]
        self.count -= 1

def numIslands(grid):
    if not grid: return 0
    rows, cols = len(grid), len(grid[0])
    uf = UnionFind(rows * cols)
    water = rows * cols          # dummy id for water cells

    for r in range(rows):
        for c in range(cols):
            if grid[r][c] == '0':
                uf.parent[r*cols + c] = water   # mark as water
                continue
            # check right and down neighbours to avoid double work
            if r+1 < rows and grid[r+1][c] == '1':
                uf.union(r*cols + c, (r+1)*cols + c)
            if c+1 < cols and grid[r][c+1] == '1':
                uf.union(r*cols + c, r*cols + (c+1))

    # count distinct roots that are not water
    roots = set()
    for r in range(rows):
        for c in range(cols):
            if grid[r][c] == '1':
                roots.add(uf.find(r*cols + c))
    return len(roots)

Why this feels like a power‑up:

• Each cell is processed once → O(m·n).
• Every union/find is practically O(1) thanks to path compression + union by size.
• No recursion depth issues, no repeated scans.

Another Real‑World LeetCode Example: Friend Circles (LeetCode 547)

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

Same idea: treat each person as a node, union whenever a direct friendship exists, and the final count tells you how many separate circles exist. The interviewer loves this because it shows you can turn a graph connectivity problem into a few lines of clean, reusable code.

Why This New Power Matters

Once you internalize Union‑Find, a whole class of problems becomes trivial:

• Detecting cycles in an undirected graph.
• Kruskal’s minimum spanning tree algorithm.
• Dynamic connectivity (e.g., “are these two computers still connected after adding/removing edges?”).
• Even grid‑based puzzles like “number of islands” or “surrounded regions”.

You stop writing ad‑hoc DFS/BFS scaffolding for every connectivity query and start thinking in terms of sets and merges. It’s a mental shift that makes your code shorter, faster, and far easier to reason about—exactly the kind of elegance that earns nods in interviews and praise in code reviews.

Your Turn

Grab a LeetCode problem that feels like a maze of nested loops (try Accounts Merge #721 or Redundant Connection #684). Implement Union‑Find from scratch, sprinkle in path compression, and watch the runtime drop from “maybe” to “definitely passes”.

When you finally see those green checkmarks, take a moment to smile—you’ve just added a reliable, battle‑tested tool to your developer’s arsenal. Happy coding, and may your sets always stay perfectly united! 🚀