Union-find and minimum spanning trees

Why it matters

Suppose a city wants to connect ten office buildings with fibre cable. Every pair of buildings could be linked, and each link has a cost that depends on the distance between them. The city wants every building connected to every other building, using the least total cable. That question is a minimum spanning tree problem. It appears in network design, in clustering data, in image segmentation and in many interview questions that ask for the cheapest way to connect things.

Two ideas solve it. The first is a small data structure called union-find, which keeps track of groups and answers "are these two things in the same group?" very quickly. The second is a pair of greedy algorithms, Kruskal's and Prim's, that build the tree one edge at a time. Union-find is also useful on its own, because many interview problems about groups, accounts or islands reduce to it.

You should be comfortable with graphs, with vector, with sort and with priority_queue. Greedy algorithms appear in the greedy module, so this lesson explains the greedy idea from the start rather than assuming it.

Groups and the question they answer

Imagine a class of students, each of whom joins a study group. Two students are in the same group when they are connected through a chain of friends, even if they never met directly. Over time, groups merge. When two groups are joined by a new friendship, they become one group.

Two questions come up again and again. The first is "which group is this student in?" The second is "are these two students in the same group?" The second can be answered by comparing the first answers. The third kind of operation merges two groups. Union-find is the structure built for exactly these three operations: find, check and merge.

Union-find stores each vertex with a pointer to a parent. A group is a little tree of pointers, and its root is the vertex whose parent is itself. The root is the group's name. To find the group of a vertex, follow the parent pointers up until you reach a root. To merge two groups, make one root point to the other root. That is the whole idea.

Two tricks keep the trees short, and short trees make the operations fast. The first trick is union by size. When merging, attach the smaller tree under the root of the larger tree. That keeps the depth of every tree at most about the logarithm of the group size. The second trick is path compression. During a find, every vertex on the walk is pointed straight at the root, so the next find on those vertices takes one step. With both tricks, a sequence of operations runs in almost constant time per operation. Strictly, the cost is proportional to the inverse Ackermann function, which is at most 4 for any input that fits in a computer. For everyday purposes, treat it as constant.

A worked example by hand

Start with six vertices, numbered 0 to 5. Each starts as its own group, so every vertex is its own root with size 1.

Unite 0 and 1. Find 0: it is its own root, so the root is 0. Find 1: the root is 1. The sizes are both 1, so attach 1 under 0. Now 0 is the root of size 2, and 1 points to 0.

Unite 2 and 3. Attach 3 under 2. Now 2 is the root of size 2.

Ask whether 0 and 1 are in the same group. The root of 0 is 0, and the root of 1 is 0. They match, so the answer is yes.

Ask whether 1 and 2 are in the same group. The root of 1 is 0, and the root of 2 is 2. They differ, so the answer is no.

Unite 1 and 2. The root of 1 is 0, with size 2. The root of 2 is 2, with size 2. The sizes are equal, so attach 2 under 0. Now 0 is the root of size 4, and the group {0, 1, 2, 3} is complete.

Ask whether 0 and 3 are in the same group. The root of 3 is found by walking 3 to 2 to 0, which is two steps. With path compression, 3 now points straight at 0, so the next find on 3 takes one step. The answer is yes.

Ask whether 4 and 5 are in the same group. Both are their own roots, and the roots differ, so the answer is no.

Two trees of group pointers, before and after find(3) compresses the path

The first program below runs exactly these operations. It reads a number of operations, then each operation is either a union or a query.

Type this into the Input box. The first line holds the number of vertices and the number of operations. Each operation is three numbers: 1 a b means unite, and 2 a b means ask whether a and b are together. The seven operations below repeat the steps of the worked example:

6 7
1 0 1
1 2 3
2 0 1
2 1 2
1 1 2
2 0 3
2 4 5

The program prints:

yes
no
yes
no

Read the code with the example in mind. The iota call fills parent with 0, 1, 2, ..., so each vertex starts as its own root. The find function is recursive. The line parent[x] = find(parent[x]); does the path compression. It walks up to the root, and on the way back it rewrites each pointer to the root. The unite function finds both roots, stops if they are the same, and otherwise attaches the smaller tree under the larger one. Notice that unite uses roots, not the original vertices, because only roots can be merged.

A common beginner question is why find returns parent[x] after the assignment. The return value is the root, because the recursion returns the root up the chain, and each vertex on the chain is updated to hold it. The last line of find is what gives the answer.

Counting components with union-find

Union-find is a quick way to count connected components, and it is often simpler than running a search. Start with every vertex in its own group. For each edge, unite its two ends. At the end, every group is a connected component, and the number of groups is the number of components. The number of groups equals the number of roots, which are the vertices with find(v) == v.

The second program counts components this way. It also reports the size of the largest component. It reads the same seven-vertex graph from the components lesson.

Type this into the Input box. It is the seven-vertex graph with four edges from the components lesson:

7 4
0 1
1 2
3 4
5 6

The program prints:

components: 3
largest: 3

The three roots are the groups {0, 1, 2}, {3, 4} and {5, 6}. The size of the first is 3, so the largest is 3. Compare this with the DFS version from the previous lesson. Both give the same answer. DFS needs a seen array and a recursive walk. Union-find needs a parent array and no walk at all, which makes it simpler when the edges arrive one at a time, as they do in a stream of friendship events.

Union-find shines when the graph changes. If new edges keep arriving and you must answer "are these two together?" after each one, rerunning DFS every time is slow. Union-find answers each question in almost constant time, and each new edge is one unite.

Why union by size matters

Path compression alone is not enough to keep the trees short when the merges are unlucky. Suppose a version of union always attaches the first root under the second. Call unite(0,1), then unite(1,2), then unite(2,3). The first call makes 0 point to 1. The second makes 1 point to 2, and the third makes 2 point to 3. Vertex 0 now sits at the bottom of a chain of four vertices, and finding its root takes three steps. A long run of such merges builds a chain as long as the number of vertices.

Now run the same three calls with union by size. unite(0,1) has two groups of size 1, so 1 goes under 0. unite(1,2) finds the root 0, which has size 2, and the lone vertex 2 goes under 0. unite(2,3) puts 3 under 0 too. Every vertex is now one step from the root 0. The bigger group always keeps its root, so the tree stays flat.

The reason is a counting argument. A tree gets one level deeper only when two trees of equal depth merge, and the merged tree then has at least twice as many vertices as before. So a tree of depth d has at least 2 to the power d vertices. Turned around, the depth is at most the base-2 logarithm of the size. Path compression then flattens whatever depth is left, so the finds stay fast.

Kruskal or Prim

Both algorithms give a minimum spanning tree, and the table below puts their differences in one place. Use it to decide which one to write when a problem does not say.

Both produce the same total, as the two programs in this lesson showed. Pick by the input you have, not by which one sounds more impressive.

Spanning trees and the minimum one

Now the main topic. A spanning tree of a connected graph with n vertices is a set of n - 1 edges that connects all the vertices with no cycle. It is called a tree because it has the shape of a tree: one connected piece, no loops. Many different spanning trees can exist for one graph.

If each edge has a weight, the minimum spanning tree, shortened to MST, is the spanning tree with the smallest total weight. In the cable example, each edge is a possible cable run, its weight is the cost of laying it, and the MST is the cheapest set of runs that connects every building.

Two facts make the problem friendly. First, every spanning tree has exactly n - 1 edges, so you never need to count how many to take. Second, a greedy approach works. You can build the MST one edge at a time by always taking a cheap edge that does not create a loop. The reason is the cut property: for any split of the vertices into two sides, the cheapest edge that crosses the split must be in some MST. A greedy method that always takes the cheapest safe edge therefore never makes a mistake.

We will look at two greedy algorithms. Both produce an MST, and they are usually taught together because they explain the same idea in different ways.

Kruskal's algorithm

Kruskal's algorithm works on the list of edges. It sorts all the edges by weight, from cheapest to most expensive. It then goes through them in order. For each edge, it asks whether the two ends are already in the same group. If they are, taking the edge would create a loop, so skip it. If they are not, take the edge and unite the two groups. Stop when you have taken n - 1 edges.

Union-find is the perfect helper here. The "same group" question is a find, and the merge is a unite. Kruskal's algorithm is therefore a sort followed by a sweep with union-find.

The first step of Kruskal is a sort, so step through a merge sort below. It sorts plain numbers, but Kruskal sorts edges by weight in the same way, and stable_sort is a merge sort.

Here is a worked example. The graph has five vertices and six edges: 0-1 costs 4, 0-2 costs 1, 2-1 costs 2, 1-3 costs 1, 2-3 costs 5 and 3-4 costs 3. Sorted by weight, the list is 0-2 (1), 1-3 (1), 2-1 (2), 3-4 (3), 0-1 (4) and 2-3 (5). The two edges of cost 1 are tied, and the order between them is fixed by how the sort treats ties. We will return to that.

Take 0-2. The groups {0} and {2} differ, so take it. Groups are now {0, 2}, {1}, {3}, {4}. The total is 1.

Take 1-3. The groups {1} and {3} differ, so take it. Groups are {0, 2}, {1, 3}, {4}. The total is 2.

Take 2-1. Vertex 2 is in {0, 2} and vertex 1 is in {1, 3}. They differ, so take it. Groups are now {0, 1, 2, 3} and {4}. The total is 4. We now have three edges, one fewer than the four we need, so we keep going.

Take 3-4. Vertex 3 is in the big group and vertex 4 is alone. Take it. Now there are four edges, which is n - 1, and every vertex is in one group. The total is 7.

Take 0-1. Both ends are in the same group, so skip it. It would create a loop. Take 2-3. Both ends are in the same group, so skip it. The MST is the four edges 0-2, 1-3, 2-1 and 3-4, with total 7.

The third program does exactly this. Its output shows which edges it takes and the total.

Type this into the Input box. The first line holds n and m, and each edge is written as a b w:

5 6
0 1 4
0 2 1
2 1 2
1 3 1
2 3 5
3 4 3

The program prints:

take 0-2 (1)
take 1-3 (1)
take 2-1 (2)
take 3-4 (3)
total: 7

Look at the sort. It uses stable_sort and not sort, and that choice matters for the output. sort does not promise to keep equal items in their original order. Two edges of weight 1 could come out in either order on different compilers, and the printed edges would differ, even though the total is the same. stable_sort promises to keep equal items in input order, so 0-2 comes before 1-3 on every compiler. The lesson is that when a result can be printed in more than one correct form, make the form deterministic.

The if (d.unite(e.a, e.b)) line is the whole algorithm. unite returns true only when the two ends were in different groups, so the edge is taken exactly when it does not create a loop. The loop keeps going after the tree is complete, but every later edge fails the test, so nothing else changes.

Prim's algorithm

Prim's algorithm takes a different route to the same answer. Instead of sorting all the edges and sweeping, it grows one tree from a starting vertex. At each step, it adds the cheapest edge that leaves the tree and goes to a vertex not yet in it. The tree grows like a plant, always sprouting from its cheapest available branch.

A min-heap makes Prim's algorithm fast. The heap holds edges that leave the tree, keyed by their weight. Pop gives the cheapest. If the vertex at the far end is already in the tree, the edge is out of date and you skip it. Otherwise you add the vertex, add its cost to the total, and push all of its edges that lead to vertices outside the tree.

Step through the min-heap below. Each step pops the cheapest pending edge, which is the same choice Prim's algorithm makes.

Here is the same graph, started at vertex 0. The tree begins as {0} with the heap holding the two edges out of 0: (4, to 1) and (1, to 2). Pop (1, to 2). Vertex 2 is not in the tree, so add it with cost 1. Its edges out are to 0 (already in the tree), to 1 at cost 2, and to 3 at cost 5. Push the ones that lead outside. The heap now holds (2, to 1), (4, to 1) and (5, to 3).

Pop (2, to 1). Vertex 1 is not in the tree, so add it with cost 2. Its edges out are to 0 (in the tree), to 2 (in the tree), and to 3 at cost 1. Push (1, to 3). The heap now holds (1, to 3), (4, to 1) and (5, to 3).

Pop (1, to 3). Vertex 3 is not in the tree, so add it with cost 1. Its edges out are to 1 (in the tree), to 2 (in the tree) and to 4 at cost 3. Push (3, to 4).

Pop (3, to 4). Vertex 4 is not in the tree, so add it with cost 3. Its only other edge leads to 3, which is in the tree.

Pop (4, to 1) and (5, to 3). Both lead to vertices already in the tree, so both are skipped. The tree has all five vertices, and the total is 0 for the start plus 1 plus 2 plus 1 plus 3, which is 7. The start is worth nothing, so the total is 7.

The fourth program does this. It prints the vertices in the order they join the tree, so you can check the trace above.

Type the same five-vertex graph into the Input box, using the same six lines as the Kruskal example. The program prints:

add 0 (weight 0)
add 2 (weight 1)
add 1 (weight 2)
add 3 (weight 1)
add 4 (weight 3)
total: 7

The first line is the start vertex, which is added with weight 0 because it costs nothing to start. The following lines match the trace. The total is the same as Kruskal's, 7. The tree itself is the same too in this example, which is not guaranteed. When weights tie, the two algorithms can pick different edges, but they always reach the same total.

The check if (inTree[v]) continue; is the stale-entry test. A vertex can be pushed several times, once through each neighbour that joined the tree. The first time it is popped is the cheapest, and every later pop is skipped.

Kruskal and Prim are both correct, and they run at similar speed in practice. Kruskal is simpler when the edges are given as a list. Prim is simpler when the graph is already an adjacency list and you want to grow from a source. Both are in the interview toolkit, and you should be able to write either.

Disconnected graphs

A spanning tree exists only when the graph is connected. If the graph has two or more pieces, the algorithms still produce a spanning forest, which is one tree per piece, with fewer than n - 1 edges in total. The total then does not answer the question "what is the cheapest way to connect everything?", because there is no way to connect everything.

A correct program must detect this. The fifth program runs Kruskal and then checks whether it took exactly n - 1 edges. If it did, the graph was connected and the total is the MST cost. If not, the graph is disconnected and the program says so.

Type this into the Input box. Four offices, but only two cables are laid, so offices 0 and 1 are joined and offices 2 and 3 are joined, and nothing joins the two pairs:

4 2
0 1 5
2 3 1

The program prints:

not connected: 2 of 3 edges needed

Now add a third cable that joins the two pairs. Type this:

4 3
0 1 5
2 3 1
1 2 2

The program prints:

tree cost: 8

The total is 5 + 1 + 2 = 8, which uses all three edges. Kruskal takes the edge 2-3 first, since it is cheapest, then 1-2, then 0-1. All three are taken, which is n - 1, so the program reports a tree.

A practical note: the check is taken == n - 1, and the problem is in the input. When the program reports "not connected", the right answer for the city is not a number. The answer is "impossible to connect everything with these cables", and that is what a good program says.

What the computer does with the groups

Each unite and find touches only a few entries of the parent array. The memory is one integer for each vertex in parent and one for each in size, so O(n) total. The edges are stored once, in the list that Kruskal sorts.

The time for Kruskal is dominated by the sort, which costs O(e log e). After sorting, the sweep does one find per endpoint for each edge. With union by size and path compression, those finds are almost constant, so the sweep is close to O(e). The total is O(e log e), and because e is at most n squared, log e is at most about 2 log n, which means the two forms are often quoted as O(e log n).

Prim with a heap runs in O(e log n), because each edge is pushed at most once from each end and each push or pop costs a logarithm of the heap size. Both are fast enough for graphs with hundreds of thousands of edges.

The recursion in find is the only place where depth matters. With union by size, every tree has depth at most log n, so the recursion is short. Without union by size, a chain could grow to length n, and the recursion would be deep. If you remove union by size, check whether a long chain can form before you trust the program.

Common mistakes

The first mistake is comparing parents instead of roots. The line if (parent[a] == parent[b]) is wrong, because the parent of a vertex is not necessarily its root. Two vertices in the same group can have different parents until path compression runs. Always compare the results of find, as if (find(a) == find(b)) does.

The second mistake is running the MST on a disconnected graph and expecting n - 1 edges. The result is a spanning forest with fewer edges. The fifth program shows the check that catches it.

The third mistake is sorting edges in descending order. That gives a maximum spanning tree, which connects everything with the largest total weight. Read the problem statement and check the direction of the sort. A small test with two edges of different weights will show the mistake at once.

The fourth mistake is using an MST to find shortest paths. The MST minimises the total weight of the tree, not the distance from one vertex to another. In the example graph, the MST uses edge 2-1 (weight 2), so the path from 0 to 1 in the tree costs 1 + 2 = 3. That happens to match the shortest path here, but it is a coincidence in this graph. In general the shortest path tree and the minimum spanning tree are different objects, and Dijkstra gives the shortest path tree.

The fifth mistake is forgetting that the sorted list in Kruskal must be sorted by weight, not by the vertex numbers. A comparator that compares a before w sorts by the wrong key, and the output looks plausible but is not minimal.

A debugging walkthrough

A classmate's Kruskal program returns a total that is too large on a small graph. The graph has four vertices and five edges, and the correct total is 6. How do you find the bug?

Print the edges in the order the program processes them, after the sort. Add a cerr line inside the loop that prints each edge and its weight. If the order is not by weight, the sort is wrong. A common cause is a comparator that returns x.a < y.a or that compares the pairs in the wrong order.

If the order is right, print the answer of unite for each edge. Each edge should be taken if and only if its ends were in different groups. Check the groups by printing find of each vertex after each union. A wrong answer often means the unite function compares the original vertices instead of their roots. The fix is to call find on both arguments before comparing, as the program does.

If both checks pass, the error is in the total. Check that the program adds e.w, the weight, and not e.a or e.b. Also check that the total is long long if the weights are large, so the sum cannot overflow.

Patterns that come up in interviews

Union-find appears in problems about connecting things with the fewest operations, such as accounts that share an email address, or pairs of friends that merge circles. The question "how many groups remain after these merges" is a count of roots. The question "does adding this edge create a cycle" is a find check, and it is exactly the test Kruskal uses.

MST appears in clustering problems, where you build the tree and cut its most expensive edges to get k clusters. It also appears in "minimum cost to connect all points" problems, where the points are on a grid and the cost is the Manhattan distance. Build all pairwise edges and run Kruskal, or use Prim with a heap if there are too many pairs.

When an interviewer asks whether an edge is in every MST or in some MST, think about ties. An edge is in some MST if no cheaper path connects its ends. An edge is in every MST if it is strictly cheaper than every other route between its ends. These are useful questions to ask out loud.

Practice

1. Use union-find to count the components of an undirected graph given as an edge list, without using DFS. Hint: start with n groups, subtract one each time a union actually merges two groups.

2. Change the Kruskal program so that it prints the total as a double average weight per edge taken. Hint: divide total by taken after the loop, and cast to double first.

3. Find out whether the shortest path from vertex 0 to vertex 4 in the sample graph is part of its MST. Hint: compare the path 0, 2, 1, 3, 4 with the four edges Kruskal took.

4. Make the Kruskal program print the edges in the order they were considered, including the ones that were skipped. Hint: print before the if, and add a word like skip or take to each line.

5. Change union-find so that it uses union by rank instead of union by size. Hint: rank is an upper bound on the depth of a tree, so increase the rank of the root only when two trees of equal rank merge.

6. Write a function that returns the number of edges in the maximum spanning tree of a connected graph. Hint: it is always n - 1, but find out why. Then explain why the total can differ from the minimum.

7. Given n points on a line at positions x and an edge between any two points with cost equal to the distance, find the minimum cost to connect all points. Hint: build the edges between neighbours in sorted order, which is enough for points on a line.

8. Explain in a paragraph why Prim's algorithm can start from any vertex and still get the same total. Hint: the cut property holds for every cut, so the greedy choice is safe from every starting point.

9. A city already has some cables laid, and each one is free to keep. Find the cheapest set of new cables to connect everything. Hint: start with the existing cables united in union-find at cost 0, then run Kruskal on the new cables only.

10. Implement a check that tells whether two vertices end up in the same group after Kruskal, using only the union-find, and without rerunning anything. Hint: after the sweep, find(a) == find(b) answers it for any pair.

What you now know

• Union-find keeps groups with a parent array. Find walks to the root, and unite joins two roots.
• Union by size and path compression keep the operations almost constant time.
• A spanning tree has n - 1 edges and connects every vertex. The minimum one has the smallest total weight.
• Kruskal sorts the edges and takes each one that does not create a loop, using union-find to test that.
• Prim grows one tree from a source, always taking the cheapest edge that leaves it, using a min-heap.
• Both produce the same total. Ties can change which edges are chosen, so use stable_sort when the output must be the same every time.
• A disconnected graph has no spanning tree. Check that Kruskal took n - 1 edges.
• Kruskal runs in O(e log e) time and Prim with a heap in O(e log n).