DEV Community

Timevolt
Timevolt

Posted on

Union-Find: The Avengers Assemble for Disjoint Sets

The Quest Begins (The “Why”)

I still remember the first time I stared at a LeetCode problem that asked me to count how many groups of friends were connected in a social network. My brute‑force solution was a nested loop that checked every pair, updated a visited array, and somehow managed to turn a simple‑looking question into an O(n²) nightmare. I felt like I was trying to move a mountain with a spoon—frustrating, slow, and definitely not the kind of code I wanted to show off in an interview.

That frustration sparked a question: Is there a smarter way to keep track of which elements belong to the same group without constantly re‑scanning everything? The answer turned out to be a humble data structure that feels like a superhero team‑up: Union‑Find, also known as Disjoint Set Union (DSU). Once I grasped it, the same problem collapsed from a painful slog to a few clean lines of code that ran in almost linear time.

The Revelation (The Insight)

At its core, Union‑Find is just a forest of trees where each node points to a parent. The root of a tree is the representative of its set. Two elements are in the same set iff they share the same root.

Why does this work so well?

  1. Find – To know the root** – By following parent pointers upward we can discover the representative. If we compress the path on the way (making every visited node point directly to the root), future finds become cheaper. It’s like giving every traveler a shortcut to the city hall after their first trip.

  2. Union by rank/size – When we merge two trees, we attach the smaller tree under the root of the larger one. This keeps the overall height shallow, preventing degenerate chains that would make find operations linear. Think of it as always putting the lighter box on top of the heavier stack so the pile never teeters.

With both tricks together, the amortized cost per operation is inverse Ackermann (α(n)), which grows so slowly that for any practical n it’s effectively constant. So a sequence of m operations on n elements runs in O(n + m α(n)) ≈ O(n + m). In interview land we usually say “almost O(1) per op,” which is more than enough to beat any naïve O(n²) approach.

Wielding the Power (Code & Examples)

The Struggle – Naïve Approach

def find_circles_naive(isConnected):
    n = len(isConnected)
    visited = [False] * n
    circles = 0

    def dfs(u):
        for v in range(n):
            if isConnected[u][v] and not visited[v]:
                visited[v] = True
                dfs(v)

    for i in range(n):
        if not visited[i]:
            visited[i] = True
            dfs(i)
            circles += 1
    return circles
Enter fullscreen mode Exit fullscreen mode

The double loop inside dfs makes each call potentially O(n), and we may invoke it n times → O(n²). Not ideal when n hits 10⁴ or more.

The Victory – Union‑Find Solution

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

    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:                     # already together
            return False
        # union by rank
        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
        return True

def find_circles_uf(isConnected):
    n = len(isConnected)
    uf = UnionFind(n)
    for i in range(n):
        for j in range(i+1, n):
            if isConnected[i][j]:
                uf.union(i, j)
    # count distinct roots
    return sum(1 for i in range(n) if uf.find(i) == i)
Enter fullscreen mode Exit fullscreen mode

What changed?

  • The double loop is still there, but the inner work is now almost O(1).
  • Path compression flattens the tree each time we look up a root, so subsequent finds are lightning fast.
  • Union by rank guarantees we never skinny‑out a tree into a long chain.

Common traps

  • Forgetting path compression – you’ll still get correct answers but the complexity degrades to O(log n) per op (or worse if you also skip union by rank).
  • Using recursion for find without hitting Python’s recursion limit on deep trees; the iterative version or the recursive version with compression is safe because the depth stays tiny after a few ops.
  • Miscounting the answer: after all unions, you must count unique roots, not just the number of union calls that returned True.

Another Classic – Redundant Connection

LeetCode 684 gives you an undirected graph that started as a tree and had one extra edge added. The task: return that edge.

def findRedundantConnection(edges):
    n = len(edges)
    uf = UnionFind(n)          # vertices are 1‑based in the problem
    for u, v in edges:
        if not uf.union(u-1, v-1):   # returns False if they were already connected
            return [u, v]
Enter fullscreen mode Exit fullscreen mode

The same Union‑Find class does the heavy lifting; we stop at the first edge that tries to unite two vertices already in the same set.

Why This New Power Matters

Mastering Union‑Find feels like unlocking a new spell in your developer grimoire. Suddenly, problems about connectivity, components, friend circles, or even detecting cycles in graphs become a matter of a few lines instead of a tangled mess of DFS/BFS bookkeeping.

You’ll walk into an interview, see a prompt like “Given n computers and a list of direct connections, return the number of separate networks,” and you’ll smile because you know the exact tool for the job. The runtime advantage isn’t just academic; on large inputs it’s the difference between a solution that passes and one that times out.

Plus, the concepts behind Union‑Find—representatives, path compression, union by rank—pop up in more advanced topics like Kruskal’s MST algorithm, dynamic connectivity, and even in some parallel computing tricks. Once you internalize them, you start seeing patterns everywhere.

Your Turn – A Little Challenge

Grab LeetCode 323: Number of Connected Components in an Undirected Graph. Given n nodes and an edge list, return how many connected components exist. Try solving it first with a simple DFS/BFS, then refactor to Union‑Find. Notice how the code shrinks and how the confidence grows.

When you’ve got it working, drop a comment with your runtime observations or a funny “I felt like a superhero when the test cases passed” moment. Happy coding, and may your sets always stay united!

Top comments (0)