Minimum Spanning Trees and Union-Find

1Connect everything, cheaply

Not the shortest route between two places — the cheapest way to make every place reachable from every other.

A spanning tree of a connected graph is a subset of its edges that touches every vertex and contains no cycle. A minimum spanning tree (MST) is a spanning tree whose total edge weight is as small as possible.

That last point catches almost everyone once, so here it is concretely:

#      A ---1--- B
#      |         /
#      3       1
#      |     /
#      C ---
#
# MST: edges A-B (1) and B-C (1), total weight 2.
#      In that tree the route from A to C is A -> B -> C, costing 2.
# But the graph has a direct A-C edge of weight 3... which is worse, fine.
#
# Now flip it:
#      A ---1--- B ---1--- C     plus a direct A-C edge of weight 1.9
# MST total = 2, and it drops the A-C edge.
# Shortest path A to C in the MST = 2. In the real graph = 1.9.
# The MST made every route slightly worse in exchange for a cheaper network.

2The cut property

Greedy algorithms usually need an excuse. This one has a theorem, and it is the reason both algorithms in this chapter are correct.

A cut is any way of splitting the vertices into two non-empty groups. An edge crosses the cut if its two endpoints are in different groups.

Its mirror image is occasionally handier:

Together the two properties give a complete decision rule: an edge is safe to take if it is the cheapest across some cut, and safe to discard if it is the dearest on some cycle. Kruskal and Prim are two ways of always having such a cut or cycle available.

3Kruskal's algorithm

Sort every edge by weight. Walk the list cheapest-first and take an edge unless its endpoints are already connected.

Kruskal ignores the shape of the graph completely and just works down a sorted list. The cut it is implicitly reasoning about is “the component containing u, versus everything else”: since we are taking edges in increasing weight order, the first edge that leaves that component is the cheapest one that does, so the cut property applies.

def kruskal(n, edges):
    """edges: list of (weight, u, v). Vertices are 0..n-1."""
    edges = sorted(edges)                # cheapest first — the whole strategy
    dsu = DSU(n)
    total, tree = 0, []

    for w, u, v in edges:
        if dsu.union(u, v):              # True if they were in different components
            total += w
            tree.append((w, u, v))
            if len(tree) == n - 1:       # a spanning tree is complete; stop early
                break

    return (total, tree) if len(tree) == n - 1 else None   # None => disconnected

The sort is O(E log E). Everything else in the loop is a union-find operation, and those are so nearly free that the sort dominates the whole algorithm. Which raises the obvious question: how does dsu.union decide, in almost no time, whether two vertices are already connected?

4Union-find (disjoint-set union)

Two operations: 'are these in the same group?' and 'merge these two groups'. Both, after two small tricks, cost essentially nothing.

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))     # everyone starts as their own president
        self.size = [1] * n              # how many members each president has
        self.count = n                   # number of separate components

    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

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False                 # already together — nothing to do
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra              # attach the SMALLER tree under the larger
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        self.count -= 1
        return True

    def connected(self, a, b):
        return self.find(a) == self.find(b)

Written naively — parent[find(b)] = find(a) with no further thought — the structure works but degenerates. Union a chain of n elements in the wrong order and you build a linked list of length n, making every find cost O(n). Two independent fixes rescue it.

α(n) is the inverse of a function that grows so violently that α is below 5 for every n that could be written down using all the atoms in the observable universe. It is not constant — the proof genuinely needs it — but you may treat it as constant in every practical sense, and everyone does.

5Prim's algorithm

Grow a single tree outward. Each step, add the cheapest edge leaving the tree. If that sounds like Dijkstra, it should.

Where Kruskal builds a forest that gradually coalesces, Prim keeps one tree and grows it. The cut it reasons about is explicit: the tree so far, versus everything else. The cheapest edge crossing that cut is safe by the cut property, so take it and repeat.

import heapq

def prim(graph, start=None):
    """graph: {vertex: [(neighbour, weight), ...]}, undirected."""
    if not graph:
        return 0, []
    start = start if start is not None else next(iter(graph))
    seen = {start}
    heap = [(w, start, v) for v, w in graph[start]]
    heapq.heapify(heap)
    total, tree = 0, []

    while heap and len(seen) < len(graph):
        w, u, v = heapq.heappop(heap)
        if v in seen:                    # stale entry: v joined via a cheaper edge
            continue
        seen.add(v)
        total += w
        tree.append((w, u, v))
        for nxt, nw in graph[v]:
            if nxt not in seen:
                heapq.heappush(heap, (nw, v, nxt))

    return (total, tree) if len(seen) == len(graph) else None

For a dense graph — E ≈ V², say a complete graph over points in the plane — the heap stops helping. Prim with a plain O(V) scan for the nearest outside vertex runs in O(V²), which beats O(V² log V) and needs no edge list at all. That version is worth knowing:

def prim_dense(n, weight):
    """weight(i, j) returns the cost of the edge i-j. Complete graph, no edge list."""
    INF = float("inf")
    best = [INF] * n          # cheapest known edge attaching i to the tree
    best[0] = 0
    inside = [False] * n
    total = 0

    for _ in range(n):
        u = min((i for i in range(n) if not inside[i]), key=lambda i: best[i])
        inside[u] = True
        total += best[u]
        for v in range(n):
            if not inside[v] and weight(u, v) < best[v]:
                best[v] = weight(u, v)

    return total

6Where this shows up

MST and union-find turn up in problems that mention neither trees nor sets.

Union-find on its own is the more frequently useful of the two. Reach for it any time the problem is about things becoming connected over time:

• Connected components of an undirected graph — union every edge, then count distinct roots. Cheaper to write than BFS and it handles edges arriving one at a time, which BFS does not.
• Cycle detection in an undirected graph — a union that returns False means both endpoints were already connected, so this edge closes a cycle.
• Merging equivalences — duplicate accounts sharing an email, image pixels of the same colour, “these two records are the same person”.
• Percolation and flood fill on a grid — treat each cell as an element and union neighbours as they open.
• Kruskal, which is what brought us here.

MST proper shows up wherever “connect everything cheaply” is the underlying question, and in a few places where it is not obvious:

• Network and utility design — the original motivation: cabling, piping, road building.
• Single-linkage clustering — build the MST, delete the k − 1 most expensive edges, and the k pieces left are exactly the clusters that maximise the minimum gap between groups. Kruskal is agglomerative clustering, run in a different order.
• Approximating the travelling salesman — walking around an MST twice gives a tour at most twice the optimal length, in polynomial time.
• Minimax / bottleneck paths — the route between two vertices that minimises its largest edge always lies in the MST. So one MST answers bottleneck queries for every pair.
• Image segmentation and maze generation — a random-weight MST on a grid is a maze with exactly one route between any two cells.

>_Playground

Kruskal and Prim on the same graph, a union-find with and without its two optimisations, and MST-based clustering.

✓Exercises

Checked automatically the moment you submit. Work top to bottom — each one assumes the last. Your answers are saved in this browser.