A Dive into Union-Find

I decided to write about this because 1) I wanted to help people learn about this algorithm, 2) it's surprisingly practical, and 3) the algorithm is beautiful. It is seen as an advanced topic, and for the longest time, I defaulted to DFS without ever asking whether something better existed.

Number of Provinces

I think examples are the best way of learning, so let's start with an example. The problem is LeetCode 547. You get an n x n matrix where isConnected[i][j] = 1 means city i and city j are directly connected. A province is a group of cities connected directly or indirectly. We want to return the number of provinces.

The obvious solution is DFS. Walk out from every city you haven't seen, mark everything reachable, and count how many times you had to start a new walk.

1def number_of_provinces_dfs(isConnected: List[List[int]]) -> int:
2    n = len(isConnected)
3    seen = set()
4    provinces = 0
5
6    def dfs(city):
7        for neighbor in range(n):
8            if isConnected[city][neighbor] and neighbor not in seen:
9                seen.add(neighbor)
10                dfs(neighbor)
11
12    for city in range(n):
13        if city not in seen:
14            seen.add(city)
15            dfs(city)
16            provinces += 1
17
18    return provinces

O(n²), because you read every cell of the matrix once. That's the baseline. Now let's build the other thing and come back to this.

Naive union-find

Union-find is a class that tracks a collection of disjoint sets. Two things live in it: a way to ask which set something belongs to, and a way to merge two sets together.

The representation is one list. parent[i] holds the node that i points to, and if parent[i] == i then i points at itself and is the root of its set. So the sets are stored as a forest, one tree per set, and the root is the name of that set.

You start with parent = [0, 1, 2, ...], so every node is its own root and every set has exactly one element.

find(i) answers "which set is i in?" It follows parent upward until it reaches a node pointing at itself, and returns that root. What matters is that two nodes in the same set always walk up to the same one.

union(i, j) merges the two sets. It calls find on both, and if the roots differ, it points one root at the other. Two trees become one. If the roots are already the same, the nodes were already in the same set and there's nothing to do.

Connectivity falls out of these. i and j are connected exactly when find(i) == find(j).

find makes one call per level on its way up, so it's O(h), where h is the height of the tree. union is two find calls plus a single pointer write, so it's O(h) too. Writing find recursively also costs O(h) space (you can also write find iteratively with a while loop and use no extra space at all, but the recursive version is shorter and reads closer to the definition).

Union-find never decides whether two nodes belong together, it just does what you tell it. It has no idea a graph exists, so picking which pairs to merge is entirely the caller's job. And when it does merge, which of the two roots ends up on top is arbitrary (keep that in mind).

Let's take a look at an example:

1class UnionFind:
2    def __init__(self, n):
3        self.parent = list(range(n))
4
5    def find(self, i):
6        if self.parent[i] == i:
7            return i
8        return self.find(self.parent[i])
9
10    def union(self, i, j):
11        root_i, root_j = self.find(i), self.find(j)
12        if root_i == root_j:
13            return False
14        self.parent[root_j] = root_i
15        return True

With that class, provinces becomes this:

1def number_of_provinces_union_find(isConnected: List[List[int]]) -> int:
2    n = len(isConnected)
3    uf = UnionFind(n)
4    provinces = n
5
6    for i in range(n):
7        for j in range(i + 1, n):
8            if isConnected[i][j] and uf.union(i, j):
9                provinces -= 1
10
11    return provinces

Start by assuming every city is its own province, so the count is n. Every time union merges two different groups, the count drops by one. If the two cities were already in the same province, union returns False and nothing happens. Note the i + 1 in the inner loop: the matrix is symmetric, so we only walk the upper triangle and skip the diagonal, which is n(n-1)/2 cells instead of n².

Reading the matrix is O(n²) on its own, and then every 1 you hit costs a union on top of that, which is O(h). So the bill is O(n² · h).

Worst case scenario

You may have realized that h can be bounded by n in the worst-case. This is caused by attaching the whole existing tree under a brand new node every time, and the forest collapses into a linked list. find is now O(n).

1class UnionFind:
2    def __init__(self, n):
3        self.parent = list(range(n))
4
5    def find(self, i):
6        if self.parent[i] == i:
7            return i
8        return self.find(self.parent[i])
9
10    def union(self, i, j):
11        root_i, root_j = self.find(i), self.find(j)
12        if root_i == root_j:
13            return False
14        self.parent[root_j] = root_i
15        return True

Union by rank

The solution is to keep track of how deep each tree is. This is the arbitrary choice from earlier, made deliberate. Instead of always hanging j under i, find the rank (depth) of each tree and hang the shorter one under the taller. The result is a tree no taller than the max rank of the two trees.

The height only increases when you merge two trees of equal height, and that requires doubling the nodes underneath. Doubling the nodes to buy one level is what gives h = O(log n). Here's the same worst-case sequence:

1class UnionFind:
2    def __init__(self, n):
3        self.parent = list(range(n))
4        self.rank = [0] * n
5
6    def find(self, i):
7        if self.parent[i] == i:
8            return i
9        return self.find(self.parent[i])
10
11    def union(self, i, j):
12        root_i, root_j = self.find(i), self.find(j)
13        if root_i == root_j:
14            return False
15
16        rank_i, rank_j = self.rank[root_i], self.rank[root_j]
17        if rank_i < rank_j:
18            self.parent[root_i] = root_j
19        if rank_i > rank_j:
20            self.parent[root_j] = root_i
21        if rank_i == rank_j:
22            self.parent[root_j] = root_i
23            self.rank[root_i] += 1
24        return True

Depth is now 1 instead of 7. find and union are now O(log n). Nice!

Path compression

This is another optimization we can do. While union by rank keeps trees from growing tall, path compression makes them shorter every time you touch one.

find already walks all the way to the root, so on the way back out it can point every node it passed directly at that root. One extra assignment does it:

1def find(self, i):
2    if self.parent[i] == i:
3        return i
4    self.parent[i] = self.find(self.parent[i])
5    return self.parent[i]

Instead of just returning the root, we store it in parent[i] first. That happens at every level on the way out of the recursion, so a single find flattens the entire path it walked. Every node it touched is now one step from the root, and the next find on any of them is immediate.

This changes what rank means. Nothing decrements it when a tree gets flattened, so it stops being the exact depth and becomes an upper bound on it. That's fine, because union by rank only needs to know which of two trees is the shorter one, and an upper bound still answers that. Fixing it up would cost more than it saves.

This is also why the bound is amortized rather than worst-case. The first find down a long path still pays for the whole walk.

Together, rank and compression put find and union at amortized O(α(n)), where α is the inverse Ackermann function. To be clear, that is much faster than O(log n). α grows so slowly that it's effectively a constant: at a billion nodes, log n is around 30 steps while α(n) is 4. It stays at 4 for any n you could fit in memory, so treat it as a constant. Space is O(n) for parent and rank, plus the recursion stack (which is mitigated by the compression).

So does it beat DFS?

Maybe. Let's go back to the province problem we introduced earlier. Every union is amortized O(α(n)) now, so the bill is O(n² · α(n)), and since α is a constant that's just O(n²), the matrix scan on its own. DFS is also O(n²), and it gets there with a visited set and no second data structure.

Clearly DFS is better-suited for this problem so why bother with union-find at all? The argument is in what the problem handed us: the whole graph all at once. What if we started to add more edges and do some queries in between? DFS has to start over from nothing every time, but with union-find, a new edge is one union and a question is two find calls. The structure never gets rebuilt!

Redundant Connection

Let's look at a case where union-find is optimal. You get a list of edges connecting n nodes, and there are exactly n of them. Connecting n nodes only takes n-1 edges, so one of these is surplus. We need to find the edge such that removing it results in a tree of n nodes. Every edge on the cycle qualifies, so the problem asks for the one that appears last in the list. This is taken from LeetCode 684.

A tree cannot have any cycles, so that means the original graph has a cycle (think about why this will ALWAYS be true in this situation). This boils the problem down to finding which edge was the one that closed a cycle when it showed up.

With DFS you'd build the graph up edge by edge, and before adding each one, check whether its two endpoints are already reachable from each other.

1def redundant_connection_dfs(edges: List[List[int]]) -> List[int]:
2    graph = defaultdict(set)
3
4    def connected(a, b, seen):
5        if a == b:
6            return True
7        seen.add(a)
8        return any(
9            connected(nxt, b, seen)
10            for nxt in graph[a]
11            if nxt not in seen
12        )
13
14    for a, b in edges:
15        if a in graph and b in graph and connected(a, b, set()):
16            return [a, b]
17        graph[a].add(b)
18        graph[b].add(a)
19
20    return []

Every check is its own traversal, so that's O(n²). DFS starting over from nothing again.

With union-find, though, it's a single pass. Here it is running on edges = [[1,2],[2,3],[3,4],[1,4],[1,5]]:

1def redundant_connection_union_find(edges: List[List[int]]) -> List[int]:
2    uf = UnionFind(len(edges) + 1)
3
4    for a, b in edges:
5        if not uf.union(a, b):
6            return [a, b]
7
8    return []

Pretty clean, right? union already tells you whether the two nodes were in different groups. If it returns False, they were already connected, so this edge closed a cycle, so this is your answer. O(n) for practical purposes.

One detail: the nodes here are labelled from 1, not 0, so the array gets sized to len(edges) + 1 and index 0 sits unused. That's the lone node off to the side in the walkthrough.

Knowing when to use it

Grouping things. Anything phrased as "how many groups" (provinces), "which items belong together," or "are these two in the same group" (find(i) == find(j)) is union-find. Others include cycle detection in an undirected graph, merging by a shared key, and edges arriving over time.

Union-find is the wrong tool for directed graphs, for anything needing the path or distance rather than a yes or no, and for edges that can be removed, since there's no undo.

Applications

Type inference in compilers. Hindley-Milner unification treats type variables as a disjoint set. When the checker learns that two types must be equal, it unions them, and find answers what a variable has been resolved to so far.

Entity resolution. Deduplicating customer or patient records, where each match rule (same email, same phone, same address) is a union and each final group is one real person.

Connected component labelling in images. Scan the pixels, union each one with its already-seen neighbours of the same colour, and every distinct blob falls out as its own set.

Maze generation. Start with walls everywhere, knock down a wall only when it joins two cells that aren't already connected, and stop when everything is one set. You get a maze with exactly one path between any two points, which is Kruskal's algorithm wearing a hat.

Thanks for reading!

I hope you enjoyed me rambling on (and hopefully teaching you something new)! This is my first technical blog, so please be merciful. Feel free to reach out! I'm always open to feedback or a chat. Stay tuned for more in the near future!