Union-find and minimum spanning trees in Python

Why it matters

A company has offices in several cities and wants to connect them all with cable. Laying cable costs money, so the company wants the set of connections that uses the least total length while still reaching every office. That set is a minimum spanning tree. The same idea appears in network design, in grouping data points into clusters, and in splitting an image into regions.

This lesson has two parts. The first is union-find, a small data structure that keeps track of groups and answers "are these two things in the same group?" quickly. The second is the minimum spanning tree itself, built with two standard methods. Kruskal's algorithm uses union-find and a sort. Prim's algorithm uses a heap from heapq, the same heap module used in the shortest-paths lesson.

The problem union-find solves

Imagine a list of friendships arriving one at a time. After each one, you want to answer two kinds of question quickly. "Are Ana and Ravi in the same friend group?" and "join Ana's group and Ravi's group into one." Each friendship merges two groups, and each question asks whether two people are already in the same group.

You could store each group as a list. To answer "same group?", you would search the lists, which is slow when there are many people. Union-find stores each group as a tree, with one member at the top called the root (or representative, the group's leader). Two people are in the same group exactly when they lead to the same root.

The structure supports two operations:

• find(x): follow the chain of parents from x up to the root, and return the root. This root is the leader of x's group.
• union(a, b): find the roots of a and b. If they are the same, the two are already in one group. If not, attach one root under the other, so the two groups become one.

The name comes from those two operations. The structure keeps a collection of disjoint sets, sets with no shared members. It can find which set an item is in, and it can union two sets.

The idea: teams with a captain

Think of each group as a team with a captain. Every member points to someone in the team, and the captain points to nobody. To find out which team you are on, you ask the person you point to, and they ask the person they point to, and so on, until someone says "I am the captain". The captain is the answer.

To merge two teams, you make one captain report to the other. The team that now has the captain at the top has all the members of both teams. The cost of merging is one link, not a copy of every member.

The array representation is simple. There is one list parent of length n. The entry parent[x] holds the parent of x. A root points to itself, so parent[r] == r is the test for a root. At the start every element is its own root, so parent = list(range(n)).

After three unions, the state for six elements can look like this: parent = [0, 0, 0, 2, 4, 5]. Element 0 is the root of the first team, which holds 0, 1, 2 and 3. Elements 4 and 5 are still single-member teams. Check element 3: its parent is 2, and 2's parent is 0, so the root of 3 is 0.

Two tricks that make it fast

Without tricks, the tree for a team can grow into a long chain, and find takes time proportional to the chain length. Two tricks keep the trees almost flat.

Union by size. When you merge two teams, attach the smaller team under the larger one. The bigger team's root stays the root, and the depth of the combined tree grows only when the smaller team was already as deep. The result is that every tree with s members has depth at most about log2 of s.

Path compression. When find walks up from x to the root, it can point every node it passes straight at the root. The next find on any of those nodes takes one step. The tree gets flatter over time, and the first walk pays for later ones.

With both tricks, a sequence of operations runs in almost constant time per operation. The exact bound is O(alpha(n)), where alpha is the inverse Ackermann function. That function grows so slowly that it is at most 4 for any input size that fits in a computer. For practical purposes, treat each operation as constant time. You do not need the proof to use the structure. What matters is that the trees stay shallow, so each find touches only a few nodes.

A trace with union by size and path compression

Take six elements, numbered 0 to 5, and perform four unions in order: union(0,1), union(2,3), union(0,2) and union(4,0). Then call find(3). Start with parent = [0, 1, 2, 3, 4, 5] and every size equal to 1.

• union(0,1): both roots are their own, and their sizes are equal, so 1 goes under 0. Now parent = [0, 0, 2, 3, 4, 5], and the size of 0 is 2.
• union(2,3): the same rule, so 3 goes under 2. Now parent = [0, 0, 2, 2, 4, 5], and the size of 2 is 2.
• union(0,2): the roots are 0 and 2, both of size 2. The test size[ra] < size[rb] is false, so 2 goes under 0. Now parent = [0, 0, 0, 2, 4, 5], and the size of 0 is 4.
• union(4,0): the roots are 4, of size 1, and 0, of size 4. The code sees that 4 is smaller and swaps the roles, so 4 goes under 0. Now parent = [0, 0, 0, 2, 0, 5], and the size of 0 is 5.

Now call find(3). The walk starts at 3, whose parent is 2. The parent of 2 is 0, which is a root. On the way back, path compression sets the parent of 3 to 0, and the parent of 2 is already 0. After find(3) the list is [0, 0, 0, 0, 0, 5]. The next find on 3 takes one step instead of two.

Notice what union by size did in the fourth step. Without it, the union of 4 and 0 could have attached the big team under the single element 4, and the root of the big team would have moved down. Each later find would then walk further. The size rule keeps the root of the big team at the top, so the trees stay shallow.

Program 1: a union-find class

The following program defines a DSU class, short for disjoint set union, the standard name for this structure. A class groups a list of data and the functions that work on it. Inside the class, self refers to the object being used. The program applies both tricks. It then runs a list of union operations, prints whether each one merged two groups, and prints the number of groups at the end. The operations are written into the program, so no input is needed.

Expected output:

0 1: True
2 3: True
1 3: True
0 2: False
groups: 2

Trace the program. Union 0 1 finds two roots, 0 and 1. They differ, and both have size 1, so the test self.size[ra] < self.size[rb] is false and 1 goes under 0. Union 2 3 does the same with 2 as root. Union 1 3 finds root 0 for element 1 and root 2 for element 3. The size of 0's team is 2 and the size of 2's team is 2, so 2 goes under 0, and the team now holds 0, 1, 2 and 3. Union 0 2 finds root 0 twice. The two are the same, so union returns False. The count of groups started at 5, and three successful unions reduce it to 2: the team of four and the single person 4.

The groups counter is kept by the loop, not by the class. Many problems need the count, and this is the simplest way to keep it correct: decrement it only when union returns True.

The first printed values are True and False. Python writes booleans with a capital letter. Inside if merged: the program can use the value directly, because Python treats True as a yes and False as a no.

find calls itself. That is safe here because union by size keeps the depth of every tree at about log2 of its size. Python stops a recursion that goes too deep, and the limit is 1000 calls by default. Union by size keeps the depth far below that. The common mistakes section shows what happens without it.

Detecting a redundant connection

A redundant connection is an edge that joins two vertices that are already connected by other edges. Adding it creates a cycle, because there is now a second route between the same pair. Union-find finds the first such edge in one pass: for each edge (a, b), if find(a) == find(b), the edge is redundant. Otherwise, union the two.

Program 2 does this with path halving, a shorter form of path compression. Each step in find points a node at its grandparent, which flattens the tree nearly as well. The program uses a plain function and a list, and it leaves out union by size to keep the code short. Without it, a tree can grow deeper before path halving flattens it, so an early find can take more steps.

The first three edges form a triangle, so the third one closes a cycle.

Expected output:

first redundant edge: 2 0

The function stops at the first redundant edge, which is (2, 0), since 0 and 1 were already joined through the first two edges. To see the second answer, remove the edge (2, 0) from the list. Then no edge closes a cycle, and the program prints no redundant edge.

Path halving is the loop body parent[x] = parent[parent[x]], followed by x = parent[x]. It moves x up two levels at a time while changing each node's parent to its grandparent. It is a little simpler than full path compression, and it gives the same almost-flat trees in practice.

Minimum spanning trees

A spanning tree of a connected undirected graph is a set of edges that connects every vertex and contains no cycle. A tree with n vertices always has exactly n - 1 edges. A graph can have many spanning trees. The minimum spanning tree (MST) is the one whose total edge weight is the smallest.

The MST answers the cable question. Each office is a vertex, each possible cable is a weighted edge, and the MST is the cheapest set of cables that reaches every office without any loop. A loop would waste money, because one cable could be removed and everything would still be connected.

Two facts make the MST easy to build greedily, which means picking the best local choice at each step. The first is that a tree never has a cycle, so each new edge must join two groups that are not yet connected. The second is the cut property: if you split the vertices into two sides, the cheapest edge crossing the split always belongs to some minimum spanning tree. In plain words, the cheapest connection between two groups is always safe to use. Both Kruskal's and Prim's algorithms rely on these two facts.

The cut property with an example

Split five cities, A to E, into two sides: {A, B} on one side and {C, D, E} on the other. The cables that cross the split are A-C (3), B-C (1) and B-D (4). The cheapest of them is B-C, at cost 1. The cut property says that this cheapest crossing cable belongs to some minimum spanning tree, so it is safe to take it. Kruskal takes B-C early. Prim, starting from A or B, also takes it, because it is the cheapest way out of that side.

This is why both algorithms are correct without any backtracking. Each decision is the cheapest way to cross a split that the algorithm has already made, and no later decision can undo it. You can check any step of either algorithm this way: take the split that the current tree or group makes, and confirm that the edge chosen is the cheapest one crossing it.

Kruskal's algorithm

Kruskal's algorithm works on the edges as a list. It is the easiest method to describe.

1. Sort all edges by weight, smallest first.
2. Go through them in order. For each edge, check whether its two ends are already connected, using union-find.
3. If they are not connected, take the edge and union the two groups. If they are connected, skip it, because it would make a cycle.
4. Stop when you have taken n - 1 edges. If the graph is disconnected, you run out of edges first.

The time is dominated by the sort, which is O(E log E). The union-find operations are almost constant time each, so they add little. The space is O(V + E).

Kruskal by hand on five cities

Here is a run of Kruskal's algorithm on paper, so you can see every decision. There are five cities, A to E, with six possible cables: A-B costs 2, A-C costs 3, B-C costs 1, B-D costs 4, C-D costs 5 and D-E costs 2. In the code these would be vertices 0 to 4, but the letters are easier to follow here.

Sort the cables by cost: B-C (1), then A-B and D-E (both 2, in either order), then A-C (3), B-D (4) and C-D (5).

Four cables are taken for five cities, which is n - 1, so the algorithm can stop. The total is 1 + 2 + 2 + 4 = 9. The cable C-D (5) is never looked at, because the loop stopped first. The cable A-C is skipped, not because it is expensive, but because A and C were already joined through B. This is the test that union-find performs for every cable, and it is why the algorithm never builds a loop.

Program 3 implements Kruskal on a four-office example. The edges are (0, 1, 10), (0, 2, 6), (0, 3, 5), (1, 3, 15) and (2, 3, 4). Each tuple is (u, v, w): one cable between u and v with cost w.

The call sorted(edges, key=lambda e: e[2]) sorts by the third item of each tuple, which is the weight. The lambda is a small unnamed function. Without the key, sorted compares the whole tuple, which starts with the first vertex number. The common mistakes section shows that result.

Expected output:

take 2-3 (4)
take 0-3 (5)
take 0-1 (10)
total 19

Follow it through. The sorted order of weights is 4, 5, 6, 10 and 15. The edge (2, 3) with weight 4 is taken first, because 2 and 3 are in different groups. The edge (0, 3) with weight 5 is taken next, joining 0 to the group of 2 and 3. The edge (0, 2) with weight 6 is skipped, because 0 and 2 are already in the same group through 3. The edge (0, 1) with weight 10 joins vertex 1, which was alone. Now the tree has 3 edges for 4 vertices, so the loop has done its job and the total is 19. The edge (1, 3) with weight 15 is never needed, so the loop reaches it and skips it.

The test taken == n - 1 checks whether the graph is connected. A disconnected graph gives a forest (a set of trees that do not join up), with fewer than n - 1 edges. Test it with an input that has an isolated vertex, and check that the output says not connected.

Prim's algorithm

Prim's algorithm grows one tree from a start vertex. It keeps a heap of edges leaving the tree, the edges that connect a tree vertex to a vertex outside. At each step it takes the cheapest such edge, adds the new vertex to the tree, and pushes the new vertex's outgoing edges onto the heap.

The idea is the cut property in action. The tree is one side of the cut, and everything else is the other side. The cheapest edge crossing the cut is always safe to add. Prim's algorithm finds that edge with the heap.

The steps are:

1. Mark the start vertex as in the tree. Push all its edges onto the heap, keyed by weight.
2. While the heap is not empty and the tree has fewer than n vertices, take the cheapest edge.
3. If its far end is already in the tree, skip it. Otherwise, add the far end to the tree, add its weight to the total, and push its outgoing edges to vertices not yet in the tree.
4. The total is the MST weight when the tree has n vertices.

The time is O(E log V) with a binary heap, and the space is O(V + E). Prim suits graphs stored as adjacency lists, and it does not need the edges sorted first. Kruskal suits an edge list, especially when the edges are already sorted, or when the graph is sparse. Both give an MST of the same total weight.

Program 4 runs Prim on the same four-office example, starting from vertex 0. The heap items are (weight, from, to) tuples. heapq compares tuples item by item, so the cheapest weight comes out first, and no comparator is needed. The program prints each edge as it joins the tree.

Expected output:

add 0-3 (5)
add 3-2 (4)
add 0-1 (10)
prim total: 19

Start at 0. Its edges are 10 to 1, 6 to 2 and 5 to 3. The heap returns 5 first, so 3 joins the tree. Vertex 3 adds edges to 1 (15) and to 2 (4). The heap now has 4 at the front, so the edge 3-2 with weight 4 joins 2. The edge 0-2 with weight 6 is still in the heap, but when it comes out, vertex 2 is already in the tree, so it is skipped. The next cheapest edge is 0-1 with weight 10, which adds vertex 1. The tree now has four vertices and the total is 19, the same as Kruskal.

The two algorithms can produce different edge lists, but they produce the same total weight. When several edges share a weight, more than one tree can have the minimum total. The total is always the same, and that is what problems usually ask for.

Kruskal or Prim

In most interview answers both are acceptable. Say which you pick and why, and state the complexity. Kruskal is the more common choice when the problem gives edges as a list, and Prim when it gives a grid or adjacency lists.

What about negative weights?

Both Kruskal and Prim work with negative weights, unlike Dijkstra. The MST is about choosing the cheapest set of edges that connects everything. A negative weight is just a cheaper edge, and the greedy choices still hold. Shortest paths need care with negative weights, but the minimum spanning tree does not.

Common mistakes and the messages they produce

Mistake 1: comparing parents instead of roots. The line if parent[a] == parent[b] checks whether the direct parents match, not whether the roots match. Two elements in one team can have different parents. Always compare find(a) with find(b). This mistake raises no error, and the groups it reports are wrong.

Mistake 2: uniting the vertices instead of their roots. Writing parent[a] = b attaches a vertex under another vertex that may not be a root, which can split a team in two and break the group. Always find the roots first, then join the roots. This mistake also raises no error.

Mistake 3: vertex numbers that start at 1. If the input numbers vertices from 1 and you use the numbers directly on a list of length n, the last vertex fails:

parent = list(range(5))
print(parent[5])

The output is this error line:

IndexError: list index out of range

The index 5 is one past the end of a list of length 5, which holds indexes 0 to 4. Subtract 1 when you read the input, or make the list n + 1 long.

Mistake 4: a recursive find on a long chain. Program 1 calls find recursively. Without union by size, a chain of parents can grow thousands of links long, and the recursion runs too deep. This code builds such a chain by hand, with 2000 elements, and calls the recursive version on the bottom element:

def find(parent, x):
    if parent[x] != x:
        parent[x] = find(parent, parent[x])
    return parent[x]
 
 
parent = list(range(1, 2000)) + [1999]
print(find(parent, 0))

The last line of the error message is:

RecursionError: maximum recursion depth exceeded

Python's default limit is 1000 nested calls, and sys.getrecursionlimit() returns that number. Union by size keeps every tree shallow, so Program 1 never gets near the limit. If you skip union by size, use the loop version of find from Program 2, which does not call itself.

Mistake 5: sorting tuples without a key. sorted(edges) compares whole tuples, so it sorts by the first vertex, then the second, then the weight. The weight is not the first item. Here is Kruskal on the same four-office example, with the key argument removed:

take 0-1 (10)
take 0-2 (6)
take 0-3 (5)
total 21

The total is 21, not 19. The program raises no error, so the mistake is easy to miss. Always check the order of the sorted list once, as the debugging walkthrough shows.

Mistake 6: using list.pop() in Prim. Prim needs the cheapest edge, so it needs the heap. list.pop() removes the last item in the list, not the smallest one. The program still runs, but the tree is not minimal. If you replace heapq.heappush(heap, ...) with heap.append(...), replace heapq.heappop(heap) with heap.pop(), and delete the heapq.heapify(heap) line, the output on the same four-office example is:

add 0-3 (5)
add 3-2 (4)
add 3-1 (15)
prim total: 24

The total is 24, not 19. The cure is heapq.heappop, which removes the smallest item.

Mistake 7: trusting the total of a disconnected graph. Kruskal returns a forest with fewer than n - 1 edges when the graph is not connected. The total is the weight of that forest, not of a spanning tree. Check the edge count, as Program 3 does, and report not connected.

Debugging walkthrough

Suppose Kruskal prints a total that is too large. Work through this checklist.

First, print the edges after sorting, with print(sorted(edges, key=lambda e: e[2])). If the weights are not in increasing order, the key is wrong. The key must select the weight, index 2 in Program 3, and not another item.

Second, print each decision. For each edge, print the two roots and whether they differ. The line that says take or skip must match what the roots say.

Third, check the union. After a union, find(a) and find(b) must return the same root. If they do not, the union did not change the parent of the smaller root, so look at the assignment parent[rv] = ru.

Fourth, check the count. The loop should end with taken equal to n - 1 for a connected graph. If it is smaller, the graph is disconnected, or the edge list is incomplete.

These four steps find most MST bugs. For union-find bugs in general, print the parent list after each union. A tree that never flattens, or a root that points to a non-root, shows up at once.

Variations and edge cases

Counting groups as edges arrive. Program 1 keeps a counter, and this is the standard approach for "how many groups after each step?" questions.

Maximum spanning tree. Sort the edges in decreasing order instead, with reverse=True, and the same code gives the heaviest set of connections.

Clustering. Run Kruskal and stop when the number of groups reaches k. The groups you have are k clusters, separated by the heaviest edges you did not take. This is single-linkage clustering, a real technique for grouping data points.

Equal weights. When several edges share a weight, the MST is not unique, but its total is. Problems that accept any MST are asking for the total or for one valid set of edges.

Edge cases to test. A graph with one vertex and no edges (the MST weight is 0), a graph with two vertices and no edge (not connected), a graph where every edge has the same weight, a graph with a self-loop (an edge from a vertex to itself, which union-find sees as a redundant edge and skips), and a graph with parallel edges between the same pair, where only the cheapest one is taken.

Interview questions this lesson answers

"Implement union-find." Parent and size lists, find with path compression, union by size. Explain each trick in one sentence.

"Does this graph have a cycle? Find the edge that creates one." Union-find: the first edge whose ends already share a root is the redundant one.

"Connect all the cities at the lowest total cost." Kruskal with union-find, or Prim with a heap. State the complexity and say why you chose it.

"How many connected groups are there after these merges?" Start the count at n and decrease it on each successful union.

"Why is union-find almost constant time?" The two tricks keep trees shallow. The bound is the inverse Ackermann function, which is at most 4 for practical inputs.

"Are Dijkstra and Prim related?" Both are greedy and use a heap. Prim has the same loop as Dijkstra, but the heap is keyed on the edge weight alone, not on the total distance from the start.

Real-world uses

MSTs design the cheapest network of cables, pipes or roads that connects every site. Union-find answers "are these two users in the same group?" in network tools, and in tools that merge duplicate records. Image segmentation groups neighbouring pixels of similar colour using a minimum spanning tree over the pixel grid. Clustering algorithms use Kruskal's order to find groups that are far apart, and percolation models in physics use union-find to decide when a material conducts.

Practice

Exercise 1 (easy). Do four union calls by hand on 5 elements: union(0,1), union(2,3), union(1,3), union(0,2). Which ones return false? Hint: check the last one.

Show the answer

Only union(0, 2) returns false, because 0 and 2 are already in the same group after the first three unions. The first three return true.

Exercise 2 (easy). What does find(x) return, and how do you test whether a and b are in the same group?

Show the answer

find(x) returns the root of x's group, the representative at the top of its tree. Two elements are in the same group exactly when find(a) == find(b). Comparing the direct parents is wrong, as explained in the common mistakes section.

Exercise 3 (easy). A tree has 10 vertices. How many edges does its spanning tree have?

Show the answer

Nine. Any spanning tree of n vertices has n - 1 edges: enough to connect them all, with no cycle.

Exercise 4 (medium). Run Kruskal by hand on the edges (0, 1, 1), (1, 2, 2) and (0, 2, 3). Which edge is skipped, and why?

Show the answer

Sorted order is weight 1, 2, 3. Take (0, 1) and (1, 2), which connect all three vertices with total 3. The edge (0, 2, 3) is skipped, because 0 and 2 are already connected through 1. The total is 3.

Exercise 5 (medium). Change the first_redundant_edge function in Program 2 so it returns a list of every redundant edge instead of stopping at the first. For the input with a triangle and a tail, what is returned?

Show the answer

Create found = [] before the loop, replace return (a, b) with found.append((a, b)), and return found after the loop. For the triangle-and-tail input, the result is [(2, 0)]. The edge (2, 3) is not redundant, because vertex 3 is alone and joins the group.

Exercise 6 (medium). Explain why Program 1 decrements groups only when union returns True. What would happen if it decremented on every call?

Show the answer

A call that returns False joins two elements already in the same group, so no group disappears. Decrementing on every call would count the redundant unions too, and the count would fall below the real number of groups. For the example input, there are 4 calls and 3 of them return True. Decrementing on every call would print groups: 1 instead of groups: 2.

Exercise 7 (medium). Write the steps Prim's algorithm takes from vertex 0 on the edges 0-1 (2), 1-2 (3), 0-2 (4) and 2-3 (1). Give the total.

Show the answer

Start at 0. Its edges are 0-1 (2) and 0-2 (4). Take 0-1 (2), adding 1. The next candidates are 0-2 (4) and 1-2 (3). Take 1-2 (3), adding 2. Its edges are 2-3 (1) and 0-2, which is skipped because 0 is already in the tree. Take 2-3 (1), adding 3. The total is 2 + 3 + 1 = 6.

Exercise 8 (medium). A graph has two components: {0, 1} with edge 0-1 (5) and {2, 3} with edge 2-3 (1). Run Kruskal and Prim from 0 and say what each returns.

Show the answer

Kruskal takes both edges, 2-3 (1) and 0-1 (5), so it returns a forest with 2 edges and total 6, which is not a spanning tree. The test taken == n - 1 fails, because 2 is not 3, so the program prints not connected. Prim from 0 takes 0-1 (5), then has no more edges that reach 2 or 3, so it stops with 1 edge, and it also prints not connected.

Exercise 9 (hard). Explain why the cut property guarantees that Prim's algorithm never makes a wrong choice. Use plain words and avoid the word "greedy".

Show the answer

At every step the tree is one side of a split, and the rest of the vertices are the other side. The cheapest edge that crosses the split is part of some minimum spanning tree. Prim always takes that cheapest crossing edge, so each edge it adds is safe, and after n - 1 steps it has a minimum spanning tree.

Exercise 10 (hard). A city has 50,000 junctions and 300,000 possible roads. Choose Kruskal or Prim, give the time, and explain why union-find keeps Kruskal fast at that size.

Show the answer

Either works. The roads are given as a list, so Kruskal is natural. Its time is O(E log E), with E = 300,000. The sort does about 300,000 times 18 comparisons, which is around 5.4 million, because log2 of 300,000 is about 18. Union-find keeps each of the 300,000 find and union operations almost constant time, because union by size and path compression keep the trees shallow. The sort dominates the total cost.

What you now know

• Union-find keeps disjoint groups with a parent list. find returns a group's root, and union joins two groups.
• Union by size and path compression keep the trees almost flat, so each operation is almost constant time.
• A redundant edge is one whose ends already share a root, so union-find finds the first cycle-closing edge in one pass.
• A minimum spanning tree has n - 1 edges and the smallest total weight. It is the cheapest way to connect everything.
• Kruskal's algorithm sorts the edges and uses union-find to skip cycle-making edges, in O(E log E) time.
• Prim's algorithm grows one tree with a heap of crossing edges, in O(E log V) time. Both give the same total weight.