The Fellowship of the Set: Mastering Union‑Find for Coding Interviews

The Quest Begins (The "Why")

I still remember the first time I saw a LeetCode problem that asked me to count the number of islands in a grid. My initial solution was a clumsy depth‑first search that visited every cell, marked it, and then recursively explored its four neighbours. It worked, but the moment I added a second test case — a massive 1000 × 1000 board — my runtime exploded and the interviewer’s eyebrows rose. I felt like I was trying to sail across an ocean with a rowboat when what I really needed was a sturdy ship.

That frustration sent me on a hunt for a data structure that could merge groups of elements quickly and tell me whether two elements belong to the same group in near‑constant time. The answer was Union‑Find, also known as Disjoint Set Union (DSU). Once I grasped it, the island problem (and dozens of others) turned from a nightmare into a straightforward “join‑and‑check” routine.

The Revelation (The Insight)

At its heart, Union‑Find solves a simple story: you have a collection of items, and over time you receive instructions of two kinds

1. Union(a, b) – declare that a and b are now in the same group.
2. Find(x) – return a representative (often called the “root”) of the group that x belongs to.

If two items share the same root, they’re in the same set; otherwise they’re separate.

The magic lies in how we store the parent links. Instead of keeping an explicit list for each set, we maintain a single array parent where parent[i] points to the representative of i’s set. Initially every element is its own parent (parent[i] = i).

When we union two sets, we attach the root of one tree to the root of the other. If we always attach the shallower tree under the deeper one (union by rank/size), the height of any tree stays logarithmic — actually, with path compression it becomes almost flat.

Path compression is the second trick: during Find(x), we recursively climb to the root, and on the way back we make every visited node point directly to that root. Future finds for those nodes become O(1).

Why does this work? Think of each set as a tangled rope. Union ties two ropes together at their ends. Path compression is like pulling the rope tight so that every knot slides straight to the knot at the end — next time you pull, you feel almost no resistance. The combination ensures that a sequence of m operations on n elements runs in O(m α(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, almost every element points straight to the ultimate boss, and asking “who’s your boss?” is a lightning‑fast lookup.

Wielding the Power (Code & Examples)

The naïve struggle (what not to do)

A common first attempt is to keep an explicit list for each component and merge by copying elements:

def num_islands_naive(grid):
    islands = []                     # list of sets, each set = cells of one island
    for r in range(len(grid)):
        for c in range(len(grid[0])):
            if grid[r][c] == '1':
                # try to merge with neighbours already seen
                merged = False
                for island in islands:
                    if (r-1,c) in island or (r+1,c) in island or (r,c-1) in island or (r,c+1) in island:
                        island.add((r,c))
                        merged = True
                        break
                if not merged:
                    islands.append({(r,c)})
    return len(islands)

The inner loop scans all existing islands for each new land cell — O(n²) in the worst case. Not acceptable for large inputs.

The victorious Union‑Find version

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

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

    def union(self, x, y):
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return                     # already in the same set
        # union by rank (attach shallower tree under deeper one)
        if self.rank[rx] < self.rank[ry]:
            self.parent[rx] = ry
        elif self.rank[rx] > self.rank[ry]:
            self.parent[ry] = rx
        else:
            self.parent[ry] = rx
            self.rank[rx] += 1

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
    count = sum(1 for r in range(rows) for c in range(cols) if grid[r][c] == '1')

    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':
                uf.union(idx, (r+1) * cols + c)
                count -= 1
            if c + 1 < cols and grid[r][c+1] == '1':
                uf.union(idx, r * cols + (c+1))
                count -= 1
    return count

What changed?

• We treat each cell as a node in a flat array (idx = r * cols + c).
• For each land cell we union it with its right and down neighbours (if they’re also land). Each union potentially reduces the number of distinct islands by one.
• find with path compression makes each lookup practically constant.

Common traps to avoid

1. Forgetting path compression – without it, the tree can degenerate into a line, turning find into O(n).
2. Union by rank/size omitted – always attaching one tree arbitrarily can still give linear height in the worst case; using rank/size guarantees logarithmic depth.
3. Double‑counting unions – only check two directions (right/down) or use a visited set to ensure each edge is processed once.

Another interview favourite: “Number of Connected Components in an Undirected Graph”

Given n nodes and an edge list, return how many connected components exist. The same Union‑Find skeleton works:

def count_components(n, edges):
    uf = UnionFind(n)
    for u, v in edges:
        uf.union(u, v)
    # count distinct roots
    return len({uf.find(i) for i in range(n)})

Again, the core logic is just a handful of union calls followed by a set of roots — O(n + m α(n)).

Why This New Power Matters

Mastering Union‑Find feels like unlocking a master key for a whole class of problems:

• Dynamic connectivity – add edges and ask “are these two nodes connected?” in near‑constant time.
• Grid‑based puzzles – islands, flood fill, maze traversal, percolation.
• Kruskal’s Minimum Spanning Tree – the algorithm hinges on efficiently checking whether adding an edge creates a cycle (i.e., whether its endpoints are already in the same set).

Once you see the pattern — union when you encounter a relationship, find when you need to know if two items share a group — you’ll start spotting it everywhere. The beauty is that the data structure is tiny (two integer arrays) yet incredibly powerful.

Your Next Quest

Here’s a challenge to cement the power:

Problem: You’re given a list of friendships (a, b) meaning person a knows person b. Friendship is transitive — if A knows B and B knows C, then A knows C. After processing all friendships, answer multiple queries of the form “do persons x and y belong to the same friend circle?”

Implement it with Union‑Find, then try to answer the queries in O(1) each after the initial preprocessing.

Give it a shot, share your solution in the comments, and let’s celebrate when your code runs faster than a superhero’s reflexes! 🚀