Union-Find in Python: Find, Union and Cycle Detection

Picture a stream of friendship events: Asha befriends Ravi, Ravi befriends Meera, and after every event someone asks whether Asha and Meera now sit in the same circle. Re-running BFS after each event costs O(V + E) every single time, so a thousand events turn into millions of steps. Union-Find (also called disjoint set union, or DSU) answers the same question in near-constant time and never even stores the graph. It is one of the smallest data structures you will ever write, and one of the few where two tiny optimisations decide whether it runs in microseconds or crawls.

Every code block and every number below comes from a real run, including a 3,000-graph random test against BFS.

When Union-Find is the right pick

Reach for it when at least one of these is true:

• The only question is “are these two in the same group?”, never “how do I get from one to the other?”.
• Connections arrive one at a time and you need the component count or a connectivity answer after each one.
• You must detect whether a new edge closes a cycle in an undirected graph.

If you need the actual route, BFS or DFS from the traversal guide is the tool. If the edges carry weights and you need the cheapest route, see Dijkstra with heapq.

Stage 1: the plain version, and why it is a trap

Every element starts as its own group, so parent[i] = i. A group is a tree stored in the parent array, and the root (the element that is its own parent) is the group’s label. find climbs to the root. union points one root at the other.

class NaiveDSU:
    def __init__(self, n):
        self.parent = list(range(n))

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

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra != rb:
            self.parent[ra] = rb

This is correct and it passes every small example. The trap is the input order. Call union(0, 1), union(1, 2), union(2, 3) and so on, and the trees grow into one long chain, so each find walks the whole chain.

Stage 2: path compression

While climbing, make every visited node point at its grandparent. This is called path halving, and it is the shortest form of path compression: one extra line, and the chain flattens each time you touch it.

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

Stage 3: union by size

Compression repairs trees after the damage. Union by size prevents the damage: always hang the smaller tree under the larger one, so the trees never get tall in the first place. Here is the complete class with both fixes.

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

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

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False          # already in the same group
        if self.size[ra] < self.size[rb]:
            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)

The return False inside union is not decoration. It tells the caller that the two elements were already connected, and that single boolean is what turns cycle detection into a one-line check in the problems below.

What the optimisations are worth

To measure it, the test built a worst-case chain with 1,999 calls of union(i, i + 1), then ran find on all 2,000 elements, counting every parent-pointer hop:

Roughly two million hops fall to about four thousand. With both fixes in place, each operation costs O(α(n)) amortised, where α is the inverse Ackermann function and stays below 5 for any n you can fit in memory. Treat it as constant. A full pass over E edges and V nodes is O(E · α(V)) time and O(V) space. Either optimisation alone already gives O(log n) amortised, so a half-remembered version still beats rerunning BFS.

Three problems, one class

Number of Provinces

Given an n × n matrix where is_connected[i][j] == 1 means two cities are linked, count the groups. The DSU already tracks it: components starts at n and drops by one on every union that merges two groups.

def find_circle_num(is_connected):
    n = len(is_connected)
    dsu = DSU(n)
    for i in range(n):
        for j in range(i + 1, n):
            if is_connected[i][j]:
                dsu.union(i, j)
    return dsu.components

Redundant Connection

A tree plus one extra edge contains exactly one cycle. Feed the edges in order. The first edge whose endpoints are already connected closes the cycle, so it is the answer. Nodes are numbered from 1, so allocate len(edges) + 1 slots.

def find_redundant_connection(edges):
    dsu = DSU(len(edges) + 1)
    for a, b in edges:
        if not dsu.union(a, b):
            return [a, b]

Graph Valid Tree

A graph on n nodes is a tree exactly when it has n − 1 edges and no cycle. Check the edge count first, then require every union to succeed.

def valid_tree(n, edges):
    if len(edges) != n - 1:
        return False
    dsu = DSU(n)
    return all(dsu.union(a, b) for a, b in edges)

The real output of these calls:

find_circle_num([[1,1,0],[1,1,0],[0,0,1]])                    -> 2
find_circle_num([[1,0,0],[0,1,0],[0,0,1]])                    -> 3
find_redundant_connection([[1,2],[1,3],[2,3]])                -> [2, 3]
find_redundant_connection([[1,2],[2,3],[3,4],[1,4],[1,5]])    -> [1, 4]
valid_tree(5, [[0,1],[0,2],[0,3],[1,4]])                      -> True
valid_tree(5, [[0,1],[1,2],[2,3],[1,3],[1,4]])                -> False

As a safety net, 3,000 random graphs (up to 12 nodes, up to 15 random edges, self-loops and repeats allowed) had their component counts compared with a BFS count. All 3,000 matched.

Bugs that survive small tests

1. Linking elements instead of roots. Writing self.parent[a] = b without calling find first looks fine for one union and breaks on the second. After union(0, 1) then union(0, 2), the second call overwrites the link from 0, so element 1 is cut loose:

buggy connected(1,2): False
correct connected(1,2): True

2. Recursive find on a long chain. The textbook recursive version is elegant, but before compression has flattened anything, a chain of 5,000 nodes raised RecursionError in the test. The iterative loop above has no such ceiling, which is why it is the version worth memorising.

3. Numbering mismatch. Problems that number nodes from 1 need an array of size n + 1. Allocating n gives an IndexError on the last node, and only on the test case that uses it.

Where Union-Find stops helping

It cannot undo a union, so problems that delete edges need an offline trick (process the operations in reverse) or a different structure. It stores membership, not routes, so “shortest path” is out. And it is undirected by nature: with directed edges and ordering constraints, the right model is topological sort.

How to present it in an interview

1. Say the idea in one sentence: every group is a tree in a parent array, and the root is the label.
2. Write the iterative find with path halving first, then add union by size and name both optimisations.
3. State the cost as “near-constant, formally O(α(n))”, and relate it to the amortised analysis in Big O notation.
4. Point out that union returns a boolean, then use it for the cycle question.

For practice beyond the three above, try Accounts Merge and Number of Connected Components in an Undirected Graph. Part of a longer plan? The coding interview roadmap shows where graphs fit.

FAQ

Do I need both path compression and union by size?

One alone is enough for good practical speed. Together they give the near-constant bound. In an interview, write compression first because it is one line, then mention union by size as the second optimisation.

What is the difference between Union-Find and DFS for connected components?

DFS finds all components in one O(V + E) pass over a graph you already have. Union-Find shines when the graph is built incrementally and you need an answer after every edge, because it never has to rescan.