Union-Find and Minimum Spanning Trees
A disjoint-set structure that answers "are these connected?" in nearly constant time, and the Kruskal and Prim algorithms that use it to connect a graph at minimum cost.
Connecting everything for the least money
A region has n farm fields and a list of m channels that could be dug between pairs of them, each with a cost. Water must be able to reach every field from every other, through any chain of channels. Which channels do you dig so that the total cost is as small as possible?
The answer is a set of n - 1 channels with no loop in it: any loop has an edge you could leave out and still be connected, and leaving it out saves money. A connected set of edges with no cycle that touches every vertex is a spanning tree, and the cheapest one is a minimum spanning tree (MST). The same structure shows up in cable and pipe layouts, in clustering points into groups, and as a building block in harder problems ("is this edge necessary?", "what is the lowest possible highest point on a route?").
Greedy ideas are usually suspicious, but here the obvious one works: look at channels from cheapest to most expensive and dig each one unless it would close a loop. Making that fast needs a data structure that answers "are these two fields already connected?" while connections are being added. That structure, union-find, is the first half of this lesson; it also solves many problems that have nothing to do with spanning trees. The previous lesson, Shortest Paths, used std::priority_queue with lazy deletion, and Prim's algorithm reuses that pattern.
Union-find: a forest of parent pointers
A disjoint-set union (DSU), or union-find structure, keeps a collection of elements split into groups, and supports two operations:
find(x): return a representative of the group containingx. Two elements are in the same group exactly whenfindreturns the same representative for both.unite(a, b): merge the groups ofaandb.
Each group is stored as a tree in which every element points to a parent and the root points to itself. The root is the representative. find walks parent pointers up to the root; unite finds both roots and makes one point at the other. The idea of these disjoint-set forests is due to Bernard Galler and Michael Fischer (1964).
Without care, the trees can degenerate into long chains and find costs O(n). Two heuristics prevent it:
- Union by size. Attach the root of the smaller tree under the root of the larger one. An element moves one level deeper only when its group is merged into a group at least as large, so its group size at least doubles each time. A group cannot exceed
nelements, so no element is ever more thanlog2 nlevels deep, and everyfindisO(log n). (Union by rank, which compares tree heights instead of sizes, gives the same bounds.) - Path compression. After
find(x)has found the root, point every element on the path straight at the root. The nextfindon any of them takes one step.
Here is a DSU with both, used for an original problem. A utility has n substations; link a b connects two of them, and ask a b asks whether they are now on the same grid. At the end, report how many separate grids there are and the size of the largest. unite returns false when a and b were already together, which keeps the grid count right and, as a side effect, detects online that a new link closes a loop.
8 9 n q; then q lines "link a b" or "ask a b"
link 1 2
link 3 4
ask 1 3
link 2 3
ask 1 4
link 4 1
link 6 7
ask 5 6
ask 7 6#include <algorithm>
#include <iostream>
#include <numeric>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
struct Dsu {
std::vector<int> parent, setSize;
explicit Dsu(int n) : parent(n), setSize(n, 1) {
std::iota(parent.begin(), parent.end(), 0); // every element is its own root
}
int find(int x) {
int root = x;
while (parent[root] != root) root = parent[root];
while (parent[x] != root) { // path compression: point the path at the root
int next = parent[x];
parent[x] = root;
x = next;
}
return root;
}
bool unite(int a, int b) { // false if a and b were already together
a = find(a);
b = find(b);
if (a == b) return false;
if (setSize[a] < setSize[b]) std::swap(a, b); // union by size
parent[b] = a;
setSize[a] += setSize[b];
return true;
}
};
void solve(std::istream& in, std::ostream& out) {
int n, q;
in >> n >> q;
Dsu dsu(n + 1); // substations 1..n
int grids = n;
while (q--) {
std::string op;
int a, b;
in >> op >> a >> b;
if (op == "link") grids -= dsu.unite(a, b);
else out << (dsu.find(a) == dsu.find(b) ? "same" : "different") << '\n';
}
int largest = 0;
for (int v = 1; v <= n; ++v) largest = std::max(largest, dsu.setSize[dsu.find(v)]);
out << grids << " grids, largest has " << largest << " substations\n";
}
int main() {
std::istringstream in(R"(8 9
link 1 2
link 3 4
ask 1 3
link 2 3
ask 1 4
link 4 1
link 6 7
ask 5 6
ask 7 6
)");
solve(in, std::cout); // on a judge: solve(std::cin, std::cout)
}Press Run to compile and run this example.
different
same
different
same
4 grids, largest has 4 substationslink 4 1 changes nothing: 1 and 4 were already connected through 2 and 3. The four grids are {1, 2, 3, 4}, {5}, {6, 7} and {8}. Both find loops are iterative on purpose. A recursive find is shorter, but on a degenerate tree it recurses as deep as the tree is tall, and without union by size that can be n levels.
How fast union-find really is
With both heuristics, any sequence of m operations on n elements takes O(m α(n)) time, where α is a functional inverse of the Ackermann function. Robert Tarjan proved this bound in 1975, and Michael Fredman and Michael Saks proved in 1989 that no disjoint-set structure can do better than α(n) amortized time per operation. Amortized means averaged over the whole sequence: a single find can still walk a long path, but it then compresses that path, so expensive operations pay for cheap later ones.
α(n) grows so slowly that it is at most 4 for any number of elements you could ever store, so in practice the cost per operation is a small constant. It is still not literally O(1), and saying so precisely is worth doing in an interview. Each heuristic alone gives O(log n) per operation: in the worst case for union by size, and amortized for path compression.
This program counts parent-pointer hops (steps of the walk to the root) for two operation sequences, with neither heuristic, with one, and with both. The chain sequence makes unite(i, 0) for every i, which, without union by size, always hangs the old tree under a new single element. The balanced sequence merges groups of equal size, like a knockout tournament, which defeats union by size on its own; it then calls find on every element twice.
#include <numeric>
#include <print>
#include <utility>
#include <vector>
// A DSU whose two heuristics can be switched off, counting parent-pointer hops.
struct CountingDsu {
bool compress, bySize;
std::vector<int> parent, setSize;
long hops = 0;
CountingDsu(int n, bool c, bool s)
: compress(c), bySize(s), parent(n), setSize(n, 1) {
std::iota(parent.begin(), parent.end(), 0);
}
int find(int x) {
int root = x;
while (parent[root] != root) root = parent[root], ++hops;
while (compress && parent[x] != root) x = std::exchange(parent[x], root);
return root;
}
void unite(int a, int b) {
a = find(a);
b = find(b);
if (a == b) return;
if (bySize && setSize[a] < setSize[b]) std::swap(a, b);
parent[b] = a; // without bySize: the root of b's set always goes under a's
setSize[a] += setSize[b];
}
};
// Chain: unite(i, 0) for i = 1, 2, ...; then find every element once.
long chain(int n, bool c, bool s) {
CountingDsu d(n, c, s);
for (int i = 1; i < n; ++i) d.unite(i, 0);
for (int i = 0; i < n; ++i) d.find(i);
return d.hops;
}
// Balanced: merge sets of equal size, like a tournament; then find every element twice.
long balanced(int n, bool c, bool s) {
CountingDsu d(n, c, s);
for (int step = 1; step < n; step *= 2)
for (int i = 0; i + step < n; i += 2 * step) d.unite(i, i + step);
for (int round = 0; round < 2; ++round)
for (int i = 0; i < n; ++i) d.find(i);
return d.hops;
}
int main() {
for (auto [name, test] : {std::pair{"chain", chain}, {"balanced", balanced}}) {
std::println("{:<9} {:>6} {:>10} {:>11} {:>8} {:>8}", name, "n", "neither",
"compression", "size", "both");
for (int n = 1024; n <= 8192; n *= 2)
std::println("{:<9} {:>6} {:>10} {:>11} {:>8} {:>8}", "", n,
test(n, false, false), test(n, true, false),
test(n, false, true), test(n, true, true));
}
}Press Run to compile and run this example.
chain n neither compression size both
1024 1046529 4088 2045 2045
2048 4190209 8184 4093 4093
4096 16769025 16376 8189 8189
8192 67092481 32760 16381 16381
balanced n neither compression size both
1024 10240 3059 10240 3059
2048 22528 6130 22528 6130
4096 49152 12273 49152 12273
8192 106496 24560 106496 24560Without heuristics the chain sequence is quadratic: each doubling of n multiplies the hops by four, and n = 8192 already needs 67 million. Union by size makes the same sequence linear. The balanced sequence shows the log n of union by size alone: in these trees the average depth is log2(n) / 2, so two rounds of find cost exactly n log2 n hops, 10240 for n = 1024, and the count grows slightly faster than n. With path compression, each element pays for its long walk once; after that, every find takes at most one hop. Use both; together they cost three extra lines.
Kruskal's algorithm
Joseph Kruskal published the "cheapest edge first" algorithm in 1956. Sort the edges by weight. Go through them in order and keep an edge if its endpoints are in different components of the edges kept so far; that is exactly dsu.unite(u, v) returning true. Stop when n - 1 edges are kept. If the list runs out first, the graph is not connected and has no spanning tree.
Back to the irrigation problem. Fields are numbered from 1, channels from 1 in input order, and the output lists the channels in the order they were chosen.
6 9 n m; then m lines "a b cost"
1 2 4
1 3 1
2 3 3
2 4 2
3 4 5
4 5 7
3 5 7
5 6 2
4 6 6#include <algorithm>
#include <iostream>
#include <numeric>
#include <sstream>
#include <tuple>
#include <utility>
#include <vector>
using ll = long long;
struct Dsu {
std::vector<int> parent, setSize;
explicit Dsu(int n) : parent(n), setSize(n, 1) {
std::iota(parent.begin(), parent.end(), 0);
}
int find(int x) {
int root = x;
while (parent[root] != root) root = parent[root];
while (parent[x] != root) x = std::exchange(parent[x], root);
return root;
}
bool unite(int a, int b) {
a = find(a);
b = find(b);
if (a == b) return false;
if (setSize[a] < setSize[b]) std::swap(a, b);
parent[b] = a;
setSize[a] += setSize[b];
return true;
}
};
struct Edge { int u, v; ll w; int id; };
void solve(std::istream& in, std::ostream& out) {
int n, m;
in >> n >> m;
std::vector<Edge> edges(m);
for (int i = 0; i < m; ++i) {
in >> edges[i].u >> edges[i].v >> edges[i].w;
edges[i].id = i + 1;
}
// Cheapest first; equal costs in input order, so the output never depends
// on how a particular std::sort orders ties.
std::ranges::sort(edges, {}, [](const Edge& e) { return std::tie(e.w, e.id); });
Dsu dsu(n + 1);
ll total = 0;
std::vector<int> chosen;
for (const Edge& e : edges) {
if (!dsu.unite(e.u, e.v)) continue; // would close a cycle
total += e.w;
chosen.push_back(e.id);
}
if (int(chosen.size()) != n - 1) { out << "impossible\n"; return; }
out << "total " << total << ", channels";
for (int id : chosen) out << ' ' << id;
out << '\n';
}
int main() {
std::istringstream in(R"(6 9
1 2 4
1 3 1
2 3 3
2 4 2
3 4 5
4 5 7
3 5 7
5 6 2
4 6 6
)");
solve(in, std::cout); // on a judge: solve(std::cin, std::cout)
}Press Run to compile and run this example.
total 14, channels 2 4 8 3 9Trace it: channel 2 (cost 1), then 4 and 8 (cost 2 each), then 3 (cost 3). Channel 1 (cost 4) would connect 1 and 2, which are already joined through 3, so it is skipped, and so is channel 5. Channel 9 connects the group {5, 6} to the rest, and the two cost-7 channels would both close loops.
Sorting costs O(E log E), which is O(E log V) for a graph without parallel edges, because then E <= V^2. The DSU work is O(E α(V)), so sorting dominates. The sort key is (cost, id), not just cost: std::sort is not stable, and two standard libraries can order equal costs differently. The total cost would not change, but the printed list of channels would. The pitfalls come back to this.
Why the greedy choice is safe: cuts and cycles
Two properties explain both algorithms in this lesson. A cut splits the vertices into two non-empty sides; an edge crosses the cut if its endpoints are on different sides.
- Cut property. For any cut, an edge of minimum weight among those crossing it belongs to some MST. If it is strictly lighter than every other crossing edge, it belongs to every MST.
- Cycle property. If an edge is strictly heavier than every other edge on some cycle, it belongs to no MST.
The cut property has a short proof by exchange. Let e be a lightest edge crossing the cut and T an MST without e. Adding e to T creates a cycle. The cycle crosses the cut at e, so it must cross it again at some other edge f of T, and w(f) >= w(e). Replacing f by e gives a spanning tree that is no heavier, so it is also an MST, and it contains e.
Kruskal follows from it. When Kruskal keeps the edge e from component C, put C on one side of a cut. Every lighter edge leaving C was already considered and would have merged something into C, so none exists, and e is a lightest edge crossing this cut. A slightly stronger form of the same exchange argument (the swapped-out edge f is never one of the edges already kept, since those do not cross the cut) shows that the kept edges always remain part of some MST.
With equal weights, "some MST" matters. A graph can have several MSTs with the same total weight and different edge sets. Kruskal and Prim may return different trees; both are correct. Compare totals, not edge lists, when you test them.
Prim's algorithm
Prim's algorithm grows one tree from a start vertex. At each step it adds the cheapest edge leaving the tree, which is safe by the cut property with the tree on one side. It was described by Vojtěch Jarník in 1930 and independently by Robert Prim in 1957 and Edsger Dijkstra in 1959 (the first of the two problems in his shortest-path paper).
The implementation looks like Dijkstra with lazy deletion. A min-heap holds candidate edges (cost, vertex, id); when an edge comes out whose vertex is already in the tree, it is stale and skipped. The difference is the key: Prim's key is the edge weight alone, while Dijkstra's is the distance from the source. Here it is on the same input.
#include <functional>
#include <iostream>
#include <queue>
#include <sstream>
#include <tuple>
#include <vector>
using ll = long long;
void solve(std::istream& in, std::ostream& out) {
int n, m;
in >> n >> m;
// adj[u] holds (cost, neighbor, channel id)
std::vector<std::vector<std::tuple<ll, int, int>>> adj(n + 1);
for (int id = 1; id <= m; ++id) {
int u, v;
ll w;
in >> u >> v >> w;
adj[u].push_back({w, v, id});
adj[v].push_back({w, u, id});
}
std::vector<bool> inTree(n + 1, false);
std::priority_queue<std::tuple<ll, int, int>, std::vector<std::tuple<ll, int, int>>,
std::greater<>> pq; // (cost, vertex, channel id)
auto addVertex = [&](int u) {
inTree[u] = true;
for (auto [w, v, id] : adj[u])
if (!inTree[v]) pq.push({w, v, id});
};
addVertex(1);
ll total = 0;
int joined = 1;
std::vector<int> chosen;
while (!pq.empty()) {
auto [w, v, id] = pq.top();
pq.pop();
if (inTree[v]) continue; // stale: v joined through a cheaper edge
total += w;
chosen.push_back(id);
++joined;
addVertex(v);
}
if (joined != n) { out << "impossible\n"; return; }
out << "total " << total << ", channels";
for (int id : chosen) out << ' ' << id;
out << '\n';
}
int main() {
std::istringstream in(R"(6 9
1 2 4
1 3 1
2 3 3
2 4 2
3 4 5
4 5 7
3 5 7
5 6 2
4 6 6
)");
solve(in, std::cout); // on a judge: solve(std::cin, std::cout)
}Press Run to compile and run this example.
total 14, channels 2 3 4 9 8Same total, 14, and here even the same set of channels, chosen in a different order because the tree grows outward from field 1. Every edge is pushed at most twice (once from each end), so the cost is O(E log E) = O(E log V). For a dense graph given as a matrix, where E is close to V^2, the version without a heap that scans an array of best connection costs for the cheapest vertex runs in O(V^2) and is the better choice.
Which one to use in a contest: Kruskal is usually shorter, works directly on an edge list, and its DSU gives the components for free. Prim fits when the graph is already an adjacency list or a dense matrix.
Stress testing Kruskal against Prim
A brute force for MST is easy to trust: try every subset of exactly n - 1 edges, keep the ones that connect all vertices (then they are spanning trees), and take the cheapest. It is exponential in E, so the graphs are tiny: up to 6 vertices and 10 edges, weights from 1 to 4 so that ties are common, with self-loops, parallel edges and disconnected graphs allowed. Both algorithms return std::optional<ll>, empty when there is no spanning tree, so the comparison checks the impossible cases as well.
#include <algorithm>
#include <bit>
#include <cstdint>
#include <functional>
#include <numeric>
#include <optional>
#include <print>
#include <queue>
#include <utility>
#include <vector>
using ll = long long;
struct Lcg {
std::uint64_t s;
std::uint32_t next() {
s = s * 6364136223846793005ULL + 1442695040888963407ULL;
return std::uint32_t(s >> 33);
}
int range(int lo, int hi) { return lo + int(next() % std::uint32_t(hi - lo + 1)); }
};
struct Dsu {
std::vector<int> parent, setSize;
explicit Dsu(int n) : parent(n), setSize(n, 1) {
std::iota(parent.begin(), parent.end(), 0);
}
int find(int x) {
int root = x;
while (parent[root] != root) root = parent[root];
while (parent[x] != root) x = std::exchange(parent[x], root);
return root;
}
bool unite(int a, int b) {
a = find(a);
b = find(b);
if (a == b) return false;
if (setSize[a] < setSize[b]) std::swap(a, b);
parent[b] = a;
setSize[a] += setSize[b];
return true;
}
};
struct Edge { int u, v; ll w; };
using Total = std::optional<ll>; // empty: no spanning tree exists
Total kruskal(int n, std::vector<Edge> edges) {
std::ranges::sort(edges, {}, &Edge::w);
Dsu dsu(n);
ll total = 0;
int used = 0;
for (auto [u, v, w] : edges) if (dsu.unite(u, v)) total += w, ++used;
return used == n - 1 ? Total{total} : std::nullopt;
}
Total prim(int n, const std::vector<Edge>& edges) {
std::vector<std::vector<std::pair<ll, int>>> adj(n);
for (auto [u, v, w] : edges) adj[u].push_back({w, v}), adj[v].push_back({w, u});
std::vector<bool> in(n, false);
std::priority_queue<std::pair<ll, int>, std::vector<std::pair<ll, int>>,
std::greater<>> pq;
pq.push({0, 0});
ll total = 0;
int joined = 0;
while (!pq.empty()) {
auto [w, v] = pq.top();
pq.pop();
if (in[v]) continue;
in[v] = true, total += w, ++joined;
for (auto e : adj[v]) if (!in[e.second]) pq.push(e);
}
return joined == n ? Total{total} : std::nullopt;
}
// Every subset of exactly n - 1 edges that connects all vertices is a spanning tree.
Total brute(int n, const std::vector<Edge>& edges) {
Total best;
int m = int(edges.size());
for (std::uint32_t mask = 0; mask < (1u << m); ++mask) {
if (std::popcount(mask) != n - 1) continue;
Dsu dsu(n);
ll total = 0;
int merges = 0;
for (int i = 0; i < m; ++i) {
if (!(mask >> i & 1)) continue;
merges += dsu.unite(edges[i].u, edges[i].v);
total += edges[i].w;
}
if (merges == n - 1 && (!best || total < *best)) best = total;
}
return best;
}
int main() {
Lcg rng{5};
int cases = 4000, failures = 0, impossible = 0;
for (int t = 0; t < cases; ++t) {
int n = rng.range(1, 6), m = rng.range(n / 2, 10);
std::vector<Edge> edges(m);
for (auto& e : edges) e = {rng.range(0, n - 1), rng.range(0, n - 1), 0};
for (auto& e : edges) e.w = rng.range(1, 4);
Total want = brute(n, edges);
impossible += !want;
failures += kruskal(n, edges) != want || prim(n, edges) != want;
}
std::println("{} random graphs ({} without a spanning tree)", cases, impossible);
std::println("disagreements: {}", failures);
}Press Run to compile and run this example.
4000 random graphs (1314 without a spanning tree)
disagreements: 0About a third of the graphs have no spanning tree, which is the point: an empty answer is an answer that has to be right. To check that the harness has teeth, we broke the code three ways. Removing the sort from Kruskal gave 1128 disagreements; returning a total without checking that n - 1 edges were used gave 1314, exactly the number of graphs without a spanning tree; removing Prim's stale-entry check gave 1860, because a vertex was then added to the tree twice.