Coding Interview: How Do You Implement Union-Find and Track Connected Components?

Prompt and Scope

There are n nodes labeled from 0 through n - 1. Initially, every node is its own connected component. Implement UnionFind with these operations:

• union(a, b) merges the components containing a and b. It returns True only when two

previously distinct components are actually merged.

• connected(a, b) reports whether the two nodes currently belong to the same component.
• count() returns the current number of connected components.

This version permits n = 0, but any operation argument must be a valid label or raise IndexError. The base problem only adds connections; it does not delete edges, and calls come from one thread. For example, after starting with n = 6 and merging (0, 1), (1, 2), and (3, 4), the three components are {0, 1, 2}, {3, 4}, and {5}. Merging (2, 4) leaves two components. A later union(0, 3) must return False without decrementing the count again.

This is a general software-engineering data-structure interview problem. Its core use case is an incremental stream of connections interleaved with many connectivity and component-count queries. If all edges arrive at once and the caller needs one component count, DFS or BFS is often more direct; recognizing that distinction is part of a strong answer.

What the Interviewer Evaluates

The first signal is state selection. Union-Find does not retain the full graph. It represents each set as a parent-pointer tree whose root is the set representative and points to itself. Consequently, connected(a, b) can compare two roots instead of traversing every stored edge.

The second signal is whether union links only roots. Writing parent[a] = b directly can move an internal node under another node and corrupt the representation of its original component. The correct sequence finds root_a and root_b, confirms that they differ, and attaches the smaller tree's root to the larger tree's root. size is meaningful only at roots and controls tree growth.

The third signal is an explanation of path compression rather than a memorized template. This implementation uses path halving: as find walks upward, it changes the current node's parent to its grandparent. That new parent is still in the same tree, so connectivity is unchanged while future paths get shorter. The iterative form also avoids recursion-depth failure on a long path.

Finally, the interviewer checks the count invariant, complexity terminology, and verification. components starts at n and decreases only after two different roots merge. Duplicate unions and self-unions cannot change it. With union by size and path compression together, operations take amortized O(α(n)) over a sequence—not strict worst-case O(1) for every call.

Clarifying Questions Before Answering

• Are connections only added, or can they also be deleted? Standard Union-Find handles additions.

After an arbitrary edge deletion, the parent forest cannot reveal whether other edges still connect the endpoints; that needs an offline method or a more advanced dynamic-connectivity structure.

• Are queries interleaved online, or are all edges provided up front? Interleaved union and

connectivity queries favor Union-Find. For one component count in a static graph, adjacency-list DFS/BFS is more transparent and retains the actual edges.

• What should union return? Here it reports whether a merge happened. That boolean supports cycle

detection directly and ensures that the component count changes exactly once.

• How should invalid nodes behave? This version raises IndexError. A contest solution can omit

validation under a guaranteed-valid contract, but Python's negative indices must not silently refer to the end of the array in a public implementation.

• Must the API report component size or enumerate members? Root size can answer size in nearly

constant amortized time. Enumerating members still costs at least the output size, and this base structure does not maintain membership lists.

• Can calls be concurrent? The base implementation mutates parent inside find, so even a

connectivity query is not read-only or thread-safe. Concurrency requires a locking contract or a specialized concurrent Union-Find algorithm.

30-Second Answer Framework

“I will keep two arrays of length n: parent[x] points to a parent, and size[root] stores the root's tree size. Initially every node is its own parent and the component count is n. find walks to a root and performs path halving by pointing each visited node to its grandparent. union finds both roots; if they are equal, it returns False. Otherwise it attaches the smaller tree's root to the larger one's root, adds their sizes, decrements the component count, and returns True. Linking two roots cannot create a cycle, and path halving only changes pointers inside one set, so connectivity remains correct. Initialization is O(n); later operations are amortized O(α(n)) with O(n) space.”

Step-by-Step Deep Dive

A direct representation assigns a component label to every node. connected is a label comparison, but merging two components requires scanning the whole array and replacing every old label, making a single union cost O(n). Another naive approach uses parent-pointer trees but attaches roots without controlling their sizes. An adversarial union order can then create a long chain and degrade find.

The recommended design maintains three pieces of state:

1. parent[x] is the parent of x, and every tree root satisfies parent[root] == root.
2. size[root] is the node count of that root's component; stale values at non-roots are never read.
3. components equals the number of roots in the parent forest.

The implementation follows. _validate is short and used only by this class, so it stays next to its call site instead of becoming a separate utility module.

class UnionFind:
    def __init__(self, n: int) -> None:
        if n < 0:
            raise ValueError("n must be non-negative")
        self.parent = list(range(n))
        self.size = [1] * n
        self.components = n

    def _validate(self, x: int) -> None:
        if x < 0 or x >= len(self.parent):
            raise IndexError("node out of range")

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

    def union(self, a: int, b: int) -> bool:
        root_a = self.find(a)
        root_b = self.find(b)
        if root_a == root_b:
            return False

        if self.size[root_a] < self.size[root_b]:
            root_a, root_b = root_b, root_a

        self.parent[root_b] = root_a
        self.size[root_a] += self.size[root_b]
        self.components -= 1
        return True

    def connected(self, a: int, b: int) -> bool:
        return self.find(a) == self.find(b)

    def count(self) -> int:
        return self.components

Correctness follows in three steps. Initially, every node is the sole root of a one-node tree, so the forest has n trees and all three invariants hold. Path halving changes x's parent to its original parent's parent. That grandparent remains on the same path to the original root, so the operation cannot cross components or change the root returned by find(x).

A union modifies only two roots. Equal roots mean that the nodes are already connected, so no state changes. For different roots, pointing root_b to root_a joins two trees into one and cannot create a cycle because those roots belonged to separate trees. The new root size is the sum of the old tree sizes, and the root count decreases by exactly one. By induction, two nodes are connected if and only if find returns the same root, and count() always matches the true component count.

Union by size guarantees that whenever a node's depth increases because its whole tree is attached, the size of its new component at least doubles. Even without path compression, tree height is at most O(log n). Combined with path halving, a sequence of m finds and unions after initialization has the amortized bound O(m α(n)). The inverse Ackermann function α grows extremely slowly. “Nearly constant amortized time” is accurate interview shorthand; “strict worst-case O(1)” is not. The two arrays use O(n) space, and iterative find uses O(1) auxiliary stack space.

The state can be traced with this operation sequence:

n = 6                         count = 6
union(0, 1) -> True           count = 5
union(1, 2) -> True           count = 4
union(3, 4) -> True           count = 3
connected(0, 2) -> True
connected(0, 4) -> False
union(2, 4) -> True           count = 2
union(0, 3) -> False          count = 2

Verification needs more than one example. Cover n = 0 without a query, a self-union at n = 1, a duplicate union, two separate components joined by a bridge, an isolated node, unions presented in opposite orders, and invalid labels -1 and n. For randomized small graphs, maintain an adjacency list as an oracle. After each edge insertion, recompute connectivity and component count with BFS and compare them step by step with Union-Find. This differential test catches subtle count and root bugs.

If all m edges are known in advance and the caller asks for one component count, adjacency-list DFS/BFS uses O(n + m) time and space and states the intent clearly. Union-Find earns its place when edges arrive incrementally and queries interleave with unions, or when Kruskal's algorithm needs to test whether an undirected edge would create a cycle. The operation model—not the mere presence of a graph—drives the choice.

High-Quality Sample Answer

“I will first confirm that relationships are only added, queries interleave with additions, and union must report whether a merge occurred. That operation model fits Union-Find. If this were one component count over a static graph, I would use DFS instead.

My state is parent, size at roots, and the current number of roots in components. Every node initially points to itself. find iteratively walks toward the root and points each visited node to its grandparent, shortening the path without recursion. union gets both roots. Equal roots return False and do not change the count. Otherwise, the smaller tree's root points to the larger tree's root, their sizes are added, and the count is decremented.

The representation remains a forest. Path halving only points a node to an ancestor in the same tree, and union connects the roots of two different trees, so neither operation creates a parent cycle. Every successful union turns exactly two trees into one, which also proves the count invariant.

Constructing the arrays costs O(n). With union by size and path halving, find, connectivity, and union are amortized O(α(n)), with O(n) space. My tests emphasize self-union and duplicate union not changing the count, bridging two large components, an isolated node, the empty structure, and invalid negative labels. I would also differential-test random small cases against BFS.”

Common Mistakes

• Assigning parent[a] = b directly → a may not be a root, so the original tree can be split or

turned into an uncontrolled chain → Find both roots and link roots only.

• Always attaching the second tree to the first → an adversarial order creates a long path → **Use

root size or rank to choose the direction.**

• Returning only parent[x] from find → a parent need not be the root, so indirect connectivity

is misclassified → Follow pointers until a self-parent root.

• Decrementing components after every union call → duplicate unions and self-unions drive the

count below reality → Update it only when the roots differ.

• Updating the old root's size after swapping roots → metadata diverges from the actual tree →

Choose the final parent root first, then link and add sizes consistently.

• Claiming worst-case O(1) per operation → the guarantee is amortized over a sequence and

includes the inverse Ackermann function → Report amortized O(α(n)).

• Ignoring Python's negative indices → find(-1) accesses the last node instead of failing →

Validate both bounds in a public implementation.

• Using basic Union-Find for arbitrary edge deletion → parent pointers do not retain alternate

paths after a deletion → **Use offline deletion processing, rollback Union-Find, or a dynamic connectivity structure.**

• Treating connected as read-only → path halving writes parent, creating races under concurrent

calls → **Define synchronization before choosing a global lock, partitioning, or a concurrent algorithm.**

Follow-Ups and How to Handle Them

Follow-up 1: How can Union-Find detect a cycle in an undirected graph?

Process edges (u, v) one at a time. If union(u, v) returns False, the endpoints were already connected before the new edge, so that edge closes a cycle. A True result only joins two previously separate components. This rule applies directly to undirected graphs. Directed-cycle detection needs a method such as three-color DFS or topological sorting.

Follow-up 2: What if the caller needs to undo the most recent union?

Use rollback Union-Find. Keep union by size and push each real change's old parent, root size, and component count onto a history stack before modifying them. Undo restores those values. Path compression is normally omitted because one find mutates many entries, inflating the rollback log and complicating boundaries. Union by size alone limits height to O(log n) and works well with divide-and-conquer over an offline operation timeline.

Follow-up 3: What if relationships can be deleted arbitrarily?

Standard Union-Find cannot answer arbitrary online deletions. If the full operation sequence is known, place each edge's active interval into a segment tree over time and traverse it with rollback Union-Find; deletions that only occur at the end can also be processed backward as additions. Truly online, frequent insertion, deletion, and query require a more advanced fully dynamic connectivity structure. Whether operations are offline is therefore a problem-defining clarification.

Follow-up 4: How would you return a component's size or all of its members?

Size is already stored at the root, so size_of(x) = size[find(x)] keeps the same amortized bound. Listing members cannot be recovered from root size alone. Scanning all nodes and comparing roots costs O(n α(n)); maintaining member sets adds merge and memory costs. Scanning is usually simpler for an occasional export. Frequent enumeration may justify a different representation.

Follow-up 5: How would you handle calls from multiple threads?

The smallest correct change is one mutex around find, union, and connected, because path halving writes the parent array. It is easy to prove but serializes all operations. Only measured contention justifies a concurrent design based on atomic compare-and-swap, deterministic linking, or partitioning. Such a design must re-prove acyclic parent pointers and atomic size and component-count updates; replacing arrays with atomic variables is not sufficient by itself.