Graph traversal, topological sort and union-find

What it is

Three patterns that share a data structure and split on a single question: do you need the path, the order, or only the grouping?

BFS explores a graph in rings of increasing distance from a source, so the first time it reaches a node, it has reached it by a shortest path. DFS follows one branch to exhaustion before backtracking, which makes it the tool for questions about structure (cycles, articulation points, strongly connected components) rather than distance. Topological sort orders the nodes of a directed acyclic graph so every edge points forward, which is the answer to "what order can I do these tasks in". Union-find, also called disjoint set union or DSU, maintains a partition of elements into disjoint groups under two operations: merge two groups, and ask whether two elements are in the same group.

Don't be confused: union-find is not a graph traversal. Union-find never stores edges and cannot tell you a path. It answers exactly one question, "are these two in the same component", and it answers it in near constant time while edges are still arriving. BFS and DFS can also compute connected components, but they need the whole graph in memory first and they cost $O(V + E)$ per query unless you precompute and cache. The rule that decides it: if edges arrive incrementally and you are asked about connectivity, use union-find; if you need the actual path, you cannot.

The other confusion worth heading off: union-find is not a tree you should read as a tree. The parent pointers form a forest, but path compression rewrites that forest constantly, and the shape after a sequence of operations tells you nothing meaningful. Only the root identity matters.

The problem it solves

Take Kruskal's minimum spanning tree algorithm. You sort the edges by weight and walk them cheapest first, taking an edge only if it joins two pieces that are not already joined. The whole algorithm hinges on that test. Do it with BFS and each test costs $O(V + E)$, which over $E$ edges gives $O(E \cdot (V+E))$, quadratic or worse. Union-find turns each test into effectively constant time and the sort dominates, giving $O(E \log E)$.

The same shape recurs whenever a relation is an equivalence relation: reflexive, symmetric and transitive. "Is in the same network partition as", "is the same user account as" during identity resolution, "is the same type variable as" during type inference. Every one of those is union-find wearing a different hat.

Topological sort solves a different starvation: given a set of tasks with prerequisites, produce an order that never violates one, and prove that no such order exists when the prerequisites are circular. Without it, you either run tasks in a hand maintained order that silently rots, or you deadlock.

Mechanics

The full runnable file:

"""Graph traversal, topological sort and union-find, with the invariants stated.

Runnable: `python3 graphs_and_union_find.py`. Every number printed here appears
in the chapter, so the chapter cannot drift from the code.
"""

from collections import deque


# --- BFS: shortest path in an unweighted graph -------------------------------
# The invariant that makes BFS correct: the queue always holds nodes in
# non-decreasing distance order, so the first time you reach a node is via a
# shortest path. That is why you mark visited on ENQUEUE, not on dequeue: marking
# on dequeue lets a node enter the queue several times and the bound degrades.
def bfs_shortest(graph, start, goal):
    if start == goal:
        return [start]
    parent = {start: None}
    q = deque([start])
    while q:
        node = q.popleft()
        for nxt in graph[node]:
            if nxt in parent:            # already reached, and by a <= distance
                continue
            parent[nxt] = node
            if nxt == goal:
                path = [goal]
                while parent[path[-1]] is not None:
                    path.append(parent[path[-1]])
                return path[::-1]
            q.append(nxt)
    return None


# --- DFS: cycle detection needs three colours, not two -----------------------
# WHITE unvisited, GREY on the current recursion stack, BLACK finished. A GREY
# neighbour is a back edge and therefore a cycle. Using only "visited" finds
# cross edges too and reports cycles that are not there.
def has_cycle(graph):
    WHITE, GREY, BLACK = 0, 1, 2
    colour = {n: WHITE for n in graph}

    def walk(n):
        colour[n] = GREY
        for nxt in graph[n]:
            if colour[nxt] == GREY:      # back edge to the current stack
                return True
            if colour[nxt] == WHITE and walk(nxt):
                return True
        colour[n] = BLACK
        return False

    return any(colour[n] == WHITE and walk(n) for n in graph)


# --- Topological sort (Kahn) -------------------------------------------------
# Kahn's version is preferred in an interview because it detects the cycle for
# free: if the output is shorter than the node count, the remainder is a cycle.
def topo_sort(graph):
    indeg = {n: 0 for n in graph}
    for n in graph:
        for m in graph[n]:
            indeg[m] += 1
    q = deque([n for n in graph if indeg[n] == 0])
    order = []
    while q:
        n = q.popleft()
        order.append(n)
        for m in graph[n]:
            indeg[m] -= 1
            if indeg[m] == 0:
                q.append(m)
    if len(order) != len(graph):
        cyclic = sorted(n for n in graph if indeg[n] > 0)
        raise ValueError(f"cycle among {cyclic}")
    return order


# --- Union-find (disjoint set union) -----------------------------------------
class UnionFind:
    """Union by size + path compression.

    The two optimisations do different jobs and you need both. Union by size
    keeps the tree shallow by hanging the smaller tree off the larger. Path
    compression flattens the path you just walked. Together the amortised cost
    per operation is the inverse Ackermann function, which is below 5 for any n
    you will ever see, so it is constant in practice.
    """

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

    def find(self, x):
        # Iterative, because the recursive version blows the stack on a long
        # chain, which is exactly the input an interviewer reaches for.
        root = x
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[x] != root:                # path compression
            self.parent[x], x = root, self.parent[x]
        return root

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False                              # already together
        if self.size[ra] < self.size[rb]:             # union by size
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        self.components -= 1
        return True

    def connected(self, a, b):
        return self.find(a) == self.find(b)


def kruskal(n, edges):
    """Minimum spanning tree. The reason union-find exists in most interviews:
    sort the edges, take one if it joins two different components."""
    total, chosen, uf = 0, [], UnionFind(n)
    for w, a, b in sorted(edges):
        if uf.union(a, b):
            total += w
            chosen.append((a, b, w))
    return total, chosen, uf.components


if __name__ == "__main__":
    g = {"a": ["b", "c"], "b": ["d"], "c": ["d", "e"], "d": ["f"], "e": ["f"], "f": []}
    print("graph:", {k: v for k, v in g.items()})
    print("bfs a->f      :", bfs_shortest(g, "a", "f"))
    print("has_cycle     :", has_cycle(g))
    print("topo order    :", topo_sort(g))

    cyc = {"x": ["y"], "y": ["z"], "z": ["x"]}
    print("cyclic graph  :", cyc)
    print("has_cycle     :", has_cycle(cyc))
    try:
        topo_sort(cyc)
    except ValueError as e:
        print("topo raises   :", e)

    print()
    uf = UnionFind(10)
    for a, b in [(0, 1), (1, 2), (3, 4), (5, 6), (6, 7), (7, 8)]:
        uf.union(a, b)
    print("after unions, components:", uf.components)
    print("connected(0,2):", uf.connected(0, 2))
    print("connected(0,3):", uf.connected(0, 3))
    print("component sizes:", sorted((uf.size[uf.find(i)] for i in range(10)), reverse=True)[:4])

    edges = [(4, 0, 1), (8, 0, 7), (11, 1, 7), (8, 1, 2), (7, 7, 8),
             (1, 6, 7), (6, 8, 6), (2, 2, 8), (2, 6, 5), (4, 2, 5),
             (14, 3, 5), (9, 3, 4), (10, 5, 4), (7, 2, 3)]
    total, chosen, comps = kruskal(9, edges)
    print()
    print("kruskal MST weight:", total)
    print("edges chosen      :", len(chosen), "for 9 nodes ->", comps, "component")

Three details in that code are the ones interviewers probe.

BFS marks visited on enqueue, not on dequeue. The line is if nxt in parent: continue followed immediately by parent[nxt] = node. If you instead mark when you pop, a node with many in-edges enters the queue once per in-edge, and the queue can hold $O(E)$ entries instead of $O(V)$. The result is still correct; the memory bound is not.

Cycle detection by DFS needs three colours, not two. A node is WHITE (untouched), GREY (on the current recursion stack) or BLACK (fully explored). A cycle exists exactly when you find an edge to a GREY node, because GREY means "I am still inside this node's call, so there is a path from it back to here". If you collapse GREY and BLACK into one "visited" flag, the diamond a to b, a to c, b to d, c to d reports a cycle that does not exist, because d is visited when c reaches it. That diamond is in the test data precisely because it is the input that separates the two implementations.

Union-find needs both optimisations, and find should be iterative. Union by size keeps trees shallow by hanging the smaller tree under the larger. Path compression flattens the path you just walked so the next find on any of those nodes is one hop. With both, the amortised cost is $O(\alpha(n))$ where $\alpha$ is the inverse Ackermann function, which is under 5 for any $n$ that fits in the observable universe. With only one of them you get $O(\log n)$, which is fine but is not what the interviewer is listening for. The iterative find matters because a recursive one recurses to depth $n$ on a degenerate chain, and "build a chain of 100,000 nodes" is a normal hostile input.

Worked example

Running the file:

graph: {'a': ['b', 'c'], 'b': ['d'], 'c': ['d', 'e'], 'd': ['f'], 'e': ['f'], 'f': []}
bfs a->f      : ['a', 'b', 'd', 'f']
has_cycle     : False
topo order    : ['a', 'b', 'c', 'd', 'e', 'f']
cyclic graph  : {'x': ['y'], 'y': ['z'], 'z': ['x']}
has_cycle     : True
topo raises   : cycle among ['x', 'y', 'z']

after unions, components: 4
connected(0,2): True
connected(0,3): False
component sizes: [4, 4, 4, 4]

kruskal MST weight: 37
edges chosen      : 8 for 9 nodes -> 1 component

Read the union-find block carefully, because it demonstrates the invariant that gets misstated in interviews. Ten elements start as ten singletons. Six successful unions run, so the component count is $10 - 6 = 4$: the group ${0,1,2}$, the group ${3,4}$, the group ${5,6,7,8}$ and the singleton ${9}$. Every successful union reduces the count by exactly one, and an unsuccessful one leaves it unchanged, which is why union returns a boolean rather than nothing. Kruskal reads that boolean directly.

The component sizes line shows [4, 4, 4, 4], and the reason is worth pausing on: it is size[find(i)] for each of the ten elements, sorted descending and truncated to four. The four largest entries are the four members of the size-4 component, all reporting the same root. It is not four separate components of size 4. self.size is only meaningful at a root, and reading it at a non-root gives a stale value left over from before that node was absorbed. Candidates who keep a size array and forget this ship an off by a lot bug.

The Kruskal run takes the classic 9-node, 14-edge graph from Sedgewick and produces weight 37 using exactly 8 edges, which is $V - 1$ as any spanning tree of a connected 9-node graph must be, ending in 1 component.

Production evidence

Type inference. Hindley-Milner unification, the algorithm under the type checkers of OCaml, Haskell and the inference parts of Rust and TypeScript, maintains equivalence classes of type variables in a union-find structure. When the checker learns that 'a and 'b must be the same type, it unions them; asking "what is this variable's type" is a find to the class representative. The near constant time is what keeps whole program inference tractable.

LLVM. LLVM ships union-find as a first class container in its ADT library, llvm::EquivalenceClasses, and uses it for congruence and equivalence problems across the optimiser rather than for one special case.

Connected component labelling in image processing. The standard two pass algorithm, which is what scipy.ndimage.label and scikit-image implement, assigns provisional labels in a raster scan and uses union-find to record that two provisional labels are actually the same region, then resolves them in the second pass. This is the canonical "edges arrive incrementally" case: you cannot build the region graph first, because discovering the regions is the task.

Topological sort in build and infrastructure tooling. Terraform builds a directed graph of resources from the references between them and walks it in topological order, which is how it parallelises independent resources and why a circular reference is a plan-time error rather than a runtime hang. Build systems in the Make and Bazel lineage do the same over targets, and Airflow rejects a DAG whose tasks form a cycle. In every one of these, the cycle detection is as much the product as the ordering, which is the argument for Kahn's algorithm below.

The debate

Kahn's algorithm or DFS post-order for topological sort? Both are $O(V+E)$ and both are correct. The real tradeoff: DFS post-order is fewer lines and needs no in-degree map, but detecting a cycle requires the three-colour bookkeeping above and recovering which nodes form the cycle takes extra work. Kahn's gets cycle detection for free, because any node left with a positive in-degree when the queue empties is in or downstream of a cycle, and that set is a usable error message. Kahn's is also naturally parallel: everything at in-degree zero can run at once, which is exactly what a build system wants.

My position: default to Kahn's. In interviews and in production the question is almost never "give me an order" alone; it is "give me an order or tell me precisely why you cannot", and Kahn's answers both from the same state. Choose DFS post-order when you are already doing a DFS for another reason, or when the graph is given as a recursive structure you would have to materialise to compute in-degrees. Note the failure mode I coded around: in topo_sort the error names the offending nodes, because "cycle detected" without the node set is an error message that costs somebody an afternoon.

When union-find is the wrong choice. Three cases, and they are all disqualifying rather than merely suboptimal. First, if you need to remove an edge or split a group. Standard union-find has no split; the structure is one-way. The problems that need it (dynamic connectivity with deletions) go to link-cut trees or Euler tour trees, and the honest interview answer is to name the constraint and say the structure changes, not to try to patch a split onto DSU. Second, if you need the path or the distance, not the grouping. Third, if the "same group" relation is not transitive, for example a similarity threshold where A is near B and B is near C but A is far from C. Union-find will happily chain those into one giant component, which is the single most common way it is misapplied in production. In entity resolution that failure has a name, over-merging, and it collapses distinct customers into one record.

Follow-up Q&A

Why is the complexity called "effectively constant" rather than constant? Because it is $O(\alpha(n))$ amortised, not $O(1)$ worst case. A single find immediately after a long chain is built can cost $O(\log n)$; path compression pays that back by making every subsequent find on those nodes cheap. Amortised is the right word and using it correctly is a signal. The bound is Tarjan's, and it is tight: no structure in this model does better.

Can you get the same bound without union by size? Path compression alone gives $O(\log n)$ amortised, and union by size alone gives $O(\log n)$ worst case. You need both for the inverse Ackermann bound. Union by rank (tree height) instead of size gives the same asymptotic result; size is easier to reason about and gives you component sizes for free, which problems frequently ask for, so I use size.

How do you get the component sizes out? Read size[find(x)] and only at a root, as above. If you want a histogram of all component sizes, iterate every element, find its root, and count roots in a dictionary; that is $O(n \alpha(n))$ and is the only correct way, because the size entries at non-roots are stale.

How would you handle a graph too large for memory? Union-find over a disk or distributed set is genuinely hard, because path compression is a random write. The practical answer is to reframe: partition the edges, run union-find per partition, then merge the partition results by unioning their boundary elements, iterating until stable. That is essentially how large scale connected components jobs work. For a single machine, note that the parent array is just two integer arrays, so 100 million elements is about 800 MB with 32-bit ints, which is often enough to make the question moot.

Your BFS returns one shortest path. What if there are several and I want all of them? Change parent from a single value to a list, and append rather than skip when you reach a node at the same distance you already reached it at, which requires tracking distance explicitly instead of inferring it. Then enumerate paths backwards from the goal. Say the cost out loud: the number of shortest paths can be exponential in $V$, so "all of them" is only reasonable if the caller consumes them lazily or you are counting rather than listing. Counting is a small DP over the BFS layers and stays linear.

Give me a problem where union-find is the trick and it is not obvious. "Given a list of accounts each with a name and a set of emails, merge accounts that share any email." The naive reading is a grouping problem; the trap is that merging is transitive through a chain of shared emails. Map each email to an element, union all emails within an account, then group emails by root. Also: processing edge removals in reverse, since a problem that asks about connectivity as edges are deleted can be run backwards, turning deletions into insertions, which union-find handles.

Common misconceptions

"Path compression makes the trees permanently flat." It flattens only the path just walked. A subsequent union can hang a tree under that root and add depth again. The bound is amortised over the sequence, not a structural guarantee at any instant.

"BFS finds shortest paths in any graph." Only in unweighted graphs, or equivalently graphs where every edge has the same weight. Add weights and you need Dijkstra; add negative weights and you need Bellman-Ford. Saying "BFS gives shortest path" without the unweighted qualifier is one of the fastest ways to lose a point on a question you knew.

"A topological order is unique." Almost never. The example above admits several valid orders; ['a','b','c','d','e','f'] is the one Kahn's produces given this insertion order. If a problem's expected output is a specific order, it has additional tie-breaking rules (usually lexicographic, which you get by swapping the queue for a heap), and you should ask rather than assume.

"I'll use DFS recursion; the graph is small." The graph in the test is small. The graph in the follow-up is 100,000 nodes in a line, and Python's default recursion limit is 1,000. Write graph DFS iteratively, or say out loud that you are choosing recursion and would convert it for production.

Interview delivery note

Say this before writing anything: "Edges arrive incrementally and I only need connectivity, not the path, so this is union-find rather than a traversal. Union by size plus path compression, effectively constant per operation." That single sentence names the structure, justifies it against the alternative, and states the complexity, which is the three things the rubric is looking for.

The senior-to-staff separator on this topic is not knowing the algorithm; everyone at this level knows it. It is naming the constraint that would break your choice before being asked. A senior candidate writes a correct union-find. A staff candidate writes the same code and adds, unprompted, "note this has no split, so if edges can be removed later this structure is wrong and we would need Euler tour trees, and separately, if grouping comes from a similarity threshold rather than a true equivalence relation, transitivity will over-merge." That is the difference between someone who can implement a pattern and someone you trust to choose one.

The Staff & Team Lead Interview Reference