Minimum Spanning Tree: Kruskal and Prim

Intuition

A spanning tree of a connected undirected graph keeps all n vertices and exactly n - 1 edges, with no cycle. A minimum spanning tree (MST) is one with the smallest total edge weight: the cheapest way to wire every city to the network.

Two greedy facts make MSTs easy:

• Cut property. Split the vertices into any two groups. The lightest edge crossing between them belongs to some MST. (If an MST skipped it, adding it would close a cycle that crosses the split at another, heavier-or-equal edge; swapping them never makes the tree worse.)
• Cycle property. The heaviest edge on a cycle is never needed (ties aside), because the rest of the cycle already connects its endpoints.

Both classic algorithms are these facts applied greedily:

• Kruskal scans edges from light to heavy and keeps an edge if it joins two different components. Union-find answers "same component?" in near-constant time.
• Prim grows one tree from a start vertex and repeatedly adds the cheapest edge leaving the tree. It looks like Dijkstra, except the heap key is the edge weight, not the distance from the start.

If the graph is disconnected, there is no spanning tree, only a minimum spanning forest (one tree per component).

When to use it in interviews / OAs

• "Connect all cities / houses / points with minimum total cost", "minimum cost to repair roads", "connect all computers".
• Points in the plane with Manhattan or Euclidean cost: an implicit complete graph, so use the O(n^2) Prim.
• "Some vertices can pay a fixed cost instead of connecting" (build a well, install a generator): add a virtual vertex 0 whose edges carry those fixed costs, then take an MST.
• Minimax / bottleneck questions: "the path whose largest edge is as small as possible". The MST contains such a path for every pair.
• "Split into k groups maximizing the minimum distance between groups": run Kruskal until exactly k components remain.

Template

Union-find (disjoint set union)

import heapq

INF = float("inf")


class DSU:
    def __init__(self, n: int):
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, x: int) -> int:
        root = x
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[x] != root:  # path compression, iteratively
            self.parent[x], x = root, self.parent[x]
        return root

    def union(self, a: int, b: int) -> bool:
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False  # already connected: this edge would close a cycle
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra  # union by size
        self.size[ra] += self.size[rb]
        return True

Kruskal

def kruskal(n: int, edges: list[tuple[int, int, int]]) -> tuple[int, list[tuple[int, int, int]]] | None:
    dsu = DSU(n)
    total, used = 0, []
    for u, v, w in sorted(edges, key=lambda e: e[2]):
        if dsu.union(u, v):
            total += w
            used.append((u, v, w))
            if len(used) == n - 1:
                break
    return (total, used) if len(used) == n - 1 else None  # None: graph is disconnected

Prim with a heap (sparse graphs)

def prim(n: int, adj: list[list[tuple[int, int]]]) -> int | None:
    in_tree = [False] * n
    pq = [(0, 0)]  # (weight of the edge that reaches the vertex, vertex)
    total = taken = 0
    while pq and taken < n:
        w, u = heapq.heappop(pq)
        if in_tree[u]:
            continue  # an older, more expensive offer for u
        in_tree[u] = True
        total += w
        taken += 1
        for v, wv in adj[u]:
            if not in_tree[v]:
                heapq.heappush(pq, (wv, v))
    return total if taken == n else None

Prim in O(n^2) for complete graphs

When every pair of points is an edge, listing n^2 / 2 edges for Kruskal costs memory and a sort. Prim with a plain array is simpler and faster:

def min_cost_connect_points(points: list[tuple[int, int]]) -> int:
    n = len(points)
    best = [INF] * n  # cheapest known edge from the tree to each vertex
    best[0] = 0
    in_tree = [False] * n
    total = 0
    for _ in range(n):
        u = min((i for i in range(n) if not in_tree[i]), key=best.__getitem__)
        in_tree[u] = True
        total += best[u]
        ux, uy = points[u]
        for v in range(n):
            if not in_tree[v]:
                d = abs(ux - points[v][0]) + abs(uy - points[v][1])
                if d < best[v]:
                    best[v] = d
    return total

Complexity

• Kruskal: O(m log m) for sorting, plus near-linear union-find work. Memory O(n + m).
• Prim with a heap: O(m log n).
• Array Prim: O(n^2) time and O(n) memory, the right choice when m is about n^2.

Common pitfalls

• Disconnected input. Check that exactly n - 1 edges were taken; otherwise report "impossible" (or return a forest if that is what is asked).
• Directed graphs. MST algorithms are for undirected graphs. The directed analogue (minimum arborescence) needs a different algorithm and is beyond interview scope.
• Union-find without path compression or union by size can degrade to O(n) per operation. A recursive find can also hit Python's recursion limit on long chains; the iterative version avoids it.
• Kruskal on a complete graph materializes n^2 / 2 edges; with n = 2000 that is two million tuples to build and sort. Use array Prim.
• Confusing MST with shortest paths. The MST does not contain shortest paths between vertices; it minimizes the total, not individual routes.
• Assuming uniqueness. Equal weights can give several MSTs with the same total. Only the total (and the multiset of weights) is guaranteed.
• Negative weights are fine here, unlike in Dijkstra.

Interview variations

• Virtual vertex. "Optimize water distribution": a well at house i costs wells[i], modeled as edge (0, i); then a plain MST over n + 1 vertices.
• Critical and pseudo-critical edges. An edge is critical if excluding it raises the MST weight (or disconnects the graph); pseudo-critical if forcing it in keeps the optimal weight. Run Kruskal once per edge with "skip this edge" or "take this edge first".
• Bottleneck paths with offline queries. "Is there a path from u to v using only edges shorter than L?": sort the queries by L, sort edges, and union edges as L grows (union-find without building any tree).
• Maximum spanning tree. Sort descending; useful for "maximize the minimum capacity on the path".
• Two-player connectivity ("remove the maximum number of edges so both Alice and Bob can traverse"): Kruskal with shared edges first, then one union-find per player.
• k-clustering. Stop Kruskal when k components remain; the next edge that would merge two components is the maximum possible spacing.