Union-Find & Disjoint Set Interview Problems: Master the Pattern

Imagine you're building a social network feature that suggests friends of friends. Or you're detecting connected components in an image. Or you're solving a coding interview problem like 'Number of Islands' or 'Accounts Merge'. All these tasks share a common need: efficiently managing and querying dynamic connectivity. This is where Union-Find, also known as Disjoint Set Union (DSU), shines.

Union-Find is a data structure that tracks a set of elements partitioned into disjoint (non-overlapping) subsets. It provides two primary operations: find (determine which subset a particular element belongs to) and union (merge two subsets into one). With optimizations like path compression and union by rank, these operations run in nearly constant time, making it a powerful tool for many graph and connectivity problems.

In this guide, you'll learn the core concepts, see how to implement DSU in Python and Java, and walk through classic interview problems step by step. We'll also cover common mistakes, debugging tips, and real-world production incidents where Union-Find saved the day. By the end, you'll be ready to tackle any Union-Find problem that comes your way.

What is Union-Find?

Union-Find, also known as Disjoint Set Union (DSU), is a data structure that keeps track of a set of elements partitioned into a number of disjoint (non-overlapping) subsets. It supports two primary operations: - Find: Determine which subset a particular element belongs to. This is typically used to check if two elements are in the same subset. - Union: Merge two subsets into a single subset.

The structure is initialized with each element in its own subset. Over time, unions connect elements, forming larger sets. The find operation can be optimized with path compression, where after finding the root of an element, we update the parent pointers along the path to point directly to the root, flattening the tree. The union operation can be optimized with union by rank or union by size, where we attach the smaller tree under the root of the larger tree to keep the tree shallow.

With these optimizations, the amortized time complexity per operation is nearly O(α(n)), where α is the inverse Ackermann function, which grows extremely slowly and is effectively constant for all practical input sizes.

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])  # path compression
        return self.parent[x]

    def union(self, x, y):
        rootX = self.find(x)
        rootY = self.find(y)
        if rootX == rootY:
            return
        # union by rank
        if self.rank[rootX] < self.rank[rootY]:
            self.parent[rootX] = rootY
        elif self.rank[rootX] > self.rank[rootY]:
            self.parent[rootY] = rootX
        else:
            self.parent[rootY] = rootX
            self.rank[rootX] += 1

Classic Problem: Number of Islands

The Number of Islands problem is a classic application of Union-Find. Given a 2D grid of '1's (land) and '0's (water), count the number of islands (connected groups of '1's). You can solve this using DFS/BFS, but Union-Find is also elegant.

Approach: Treat each cell as an element. For each '1', union it with its right and down neighbors if they are also '1'. After processing all cells, count the number of distinct roots among '1's.

Time complexity: O(m n α(mn)), where m and n are grid dimensions. Space: O(m n) for the parent array.

Let's implement this in Python.

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
        self.count = 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):
        rootX = self.find(x)
        rootY = self.find(y)
        if rootX != rootY:
            if self.rank[rootX] < self.rank[rootY]:
                self.parent[rootX] = rootY
            elif self.rank[rootX] > self.rank[rootY]:
                self.parent[rootY] = rootX
            else:
                self.parent[rootY] = rootX
                self.rank[rootX] += 1
            self.count -= 1

def numIslands(grid):
    if not grid:
        return 0
    m, n = len(grid), len(grid[0])
    uf = UnionFind(m * n)
    land_cells = 0
    for i in range(m):
        for j in range(n):
            if grid[i][j] == '1':
                land_cells += 1
                idx = i * n + j
                if i + 1 < m and grid[i+1][j] == '1':
                    uf.union(idx, (i+1)*n + j)
                if j + 1 < n and grid[i][j+1] == '1':
                    uf.union(idx, i*n + j+1)
    # count components among land cells
    roots = set()
    for i in range(m):
        for j in range(n):
            if grid[i][j] == '1':
                roots.add(uf.find(i*n + j))
    return len(roots)

Problem: Accounts Merge

Given a list of accounts where each account has a name and a list of emails, merge accounts that share common emails. The output should have one account per person (name and sorted emails).

This problem is a classic Union-Find application. We treat each email as an element. For each account, we union the first email with all other emails in that account. After processing all accounts, we group emails by their root. Then for each group, we take the name from any account (they all have the same name) and output sorted emails.

Time complexity: O(N * α(N)) where N is total number of emails. Space: O(N).

def accountsMerge(accounts):
    uf = UnionFind(len(accounts))
    email_to_id = {}
    for i, account in enumerate(accounts):
        for email in account[1:]:
            if email in email_to_id:
                uf.union(i, email_to_id[email])
            else:
                email_to_id[email] = i
    # group emails by root account
    root_to_emails = defaultdict(list)
    for email, id in email_to_id.items():
        root = uf.find(id)
        root_to_emails[root].append(email)
    # build result
    result = []
    for root, emails in root_to_emails.items():
        name = accounts[root][0]
        result.append([name] + sorted(emails))
    return result

Problem: Redundant Connection

In this problem, you are given a graph that started as a tree (n nodes, n-1 edges) but one extra edge was added, creating a cycle. Find the edge that can be removed to make it a tree again. If multiple edges can be removed, return the one that appears last in the input.

Union-Find is perfect here: iterate through edges, for each edge (u, v), if u and v are already in the same set, this edge creates a cycle and is the redundant connection. Otherwise, union them.

Time complexity: O(n * α(n)). Space: O(n).

def findRedundantConnection(edges):
    n = len(edges)
    uf = UnionFind(n + 1)  # nodes are 1-indexed
    for u, v in edges:
        if uf.find(u) == uf.find(v):
            return [u, v]
        uf.union(u, v)
    return []

Problem: Longest Consecutive Sequence

Given an unsorted array of integers, find the length of the longest consecutive elements sequence. For example, [100, 4, 200, 1, 3, 2] returns 4 (sequence [1,2,3,4]).

Union-Find can solve this: treat each number as an element. For each number, if num-1 exists, union them; if num+1 exists, union them. Then find the size of the largest set.

Alternatively, use a hash set for O(n) time. But Union-Find is a valid approach.

Time complexity: O(n * α(n)). Space: O(n).

def longestConsecutive(nums):
    if not nums:
        return 0
    uf = UnionFind(len(nums))
    num_to_idx = {}
    for i, num in enumerate(nums):
        if num in num_to_idx:
            continue
        num_to_idx[num] = i
        if num - 1 in num_to_idx:
            uf.union(i, num_to_idx[num-1])
        if num + 1 in num_to_idx:
            uf.union(i, num_to_idx[num+1])
    # count sizes
    size = [0] * len(nums)
    for i in range(len(nums)):
        root = uf.find(i)
        size[root] += 1
    return max(size)

Problem: Satisfiability of Equality Equations

Given an array of equations like 'a==b' and 'a!=b', determine if all equations can be satisfied. This is a classic Union-Find problem: first process all '==' equations to union variables, then check '!=' equations: if two variables are in the same set, return false.

Time complexity: O(n * α(26)) since only lowercase letters. Space: O(26).

def equationsPossible(equations):
    uf = UnionFind(26)
    for eq in equations:
        if eq[1] == '=':
            uf.union(ord(eq[0]) - ord('a'), ord(eq[3]) - ord('a'))
    for eq in equations:
        if eq[1] == '!':
            if uf.find(ord(eq[0]) - ord('a')) == uf.find(ord(eq[3]) - ord('a')):
                return False
    return True

Advanced: Dynamic Connectivity with Offline Queries

Sometimes you need to answer connectivity queries after a series of union operations. Union-Find can handle this online. But what if you have to answer queries after some unions are undone? That's harder. However, for offline queries (all queries known in advance), you can use Union-Find with a divide-and-conquer approach (DSU with rollback). This is an advanced topic.

For interview purposes, focus on the standard Union-Find. But knowing that Union-Find can be extended to support undo operations (by storing stack of changes) can impress interviewers.

Time complexity: O((n+q) log n) for offline queries with rollback.

class DSUWithRollback:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n
        self.history = []

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

    def union(self, a, b):
        a = self.find(a)
        b = self.find(b)
        if a == b:
            return False
        if self.size[a] < self.size[b]:
            a, b = b, a
        self.history.append((b, self.parent[b], a, self.size[a]))
        self.parent[b] = a
        self.size[a] += self.size[b]
        return True

    def snapshot(self):
        return len(self.history)

    def rollback(self, snap):
        while len(self.history) > snap:
            b, parent_b, a, size_a = self.history.pop()
            self.parent[b] = parent_b
            self.size[a] = size_a