206. Prim Minimum Spanning Tree

Why this mechanism works

Grow one connected tree rather than a forest. A min-heap stores edges leaving visited vertices; the cheapest edge reaching an unseen vertex safely extends the tree across its current cut. Edges whose destination is already visited are stale candidates and must be skipped.

State and boundaries

seen is the current tree's vertex set. Heap entries are candidate (edge_weight,destination) pairs from that set, plus the initial zero-cost root entry.

Protocol

1. Build undirected weighted adjacency.
2. Seed (0,0).
3. Pop the cheapest candidate.
4. Skip an already included destination.
5. Include the new vertex, add its edge cost, and push edges to unseen neighbors.
6. Return total only if every vertex was included.

Visualization

tree {0}; frontier edges weight 1 to 1, weight 5 to 2
take 1; new edge weight 2 to 2
take 2 -> total 3; stale weight-5 entry is ignored

Python implementation

import heapq

def prim_mst(n, edges):
    if n == 0:
        return 0
    adj = [[] for _ in range(n)]
    for u, v, weight in edges:
        adj[u].append((v, weight))
        adj[v].append((u, weight))
    heap, seen, total = [(0, 0)], set(), 0
    while heap:
        weight, u = heapq.heappop(heap)
        if u in seen:
            continue
        seen.add(u)
        total += weight
        for v, cost in adj[u]:
            if v not in seen:
                heapq.heappush(heap, (cost, v))
    return total if len(seen) == n else None

Worked execution

Recorded state changes after statements; L numbers refer to the solution file. Long traces are explicitly condensed.

L4 prim_mst: if n == 0:
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]

L6 prim_mst: adj = [[] for _ in range(n)]
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; _=0

L6 prim_mst: adj = [[] for _ in range(n)]
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; _=1

L6 prim_mst: adj = [[] for _ in range(n)]
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; _=2

L6 prim_mst: adj = [[] for _ in range(n)]
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; adj=[[], [], []]

L7 prim_mst: for u, v, weight in edges:
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; adj=[[], [], []]; u=0; v=1; weight=1

L8 prim_mst: adj[u].append((v, weight))
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; adj=[[(..., ...)], [], []]; u=0; v=1; weight=1

L9 prim_mst: adj[v].append((u, weight))
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; adj=[[(..., ...)], [(..., ...)], []]; u=0; v=1; weight=1

... 21 intermediate state changes omitted ...

L16 prim_mst: total += weight
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; adj=[[(..., ...), (..., ...)], [(..., ...), (..., ...)], [(..., ...), (..., ...)]]; u=2; v=2; weight=2; heap=[(5, 2)]; seen={0, 1, 2}; total=3; cost=2

L17 prim_mst: for v, cost in adj[u]:
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; adj=[[(..., ...), (..., ...)], [(..., ...), (..., ...)], [(..., ...), (..., ...)]]; u=2; v=0; weight=2; heap=[(5, 2)]; seen={0, 1, 2}; total=3; cost=5

L17 prim_mst: for v, cost in adj[u]:
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; adj=[[(..., ...), (..., ...)], [(..., ...), (..., ...)], [(..., ...), (..., ...)]]; u=2; v=1; weight=2; heap=[(5, 2)]; seen={0, 1, 2}; total=3; cost=2

L12 prim_mst: weight, u = heapq.heappop(heap)
n=3; edges=[(0, 1, 1), (0, 2, 5), (1, 2, 2)]; adj=[[(..., ...), (..., ...)], [(..., ...), (..., ...)], [(..., ...), (..., ...)]]; u=2; v=1; weight=5; heap=[]; seen={0, 1, 2}; total=3; cost=2

L4 prim_mst: if n == 0:
n=2; edges=[]

L6 prim_mst: adj = [[] for _ in range(n)]
n=2; edges=[]; _=0

L6 prim_mst: adj = [[] for _ in range(n)]
n=2; edges=[]; _=1

L6 prim_mst: adj = [[] for _ in range(n)]
n=2; edges=[]; adj=[[], []]

L10 prim_mst: heap, seen, total = [(0, 0)], set(), 0
n=2; edges=[]; adj=[[], []]; heap=[(0, 0)]; seen=set(); total=0

L12 prim_mst: weight, u = heapq.heappop(heap)
n=2; edges=[]; adj=[[], []]; heap=[]; seen=set(); total=0; u=0; weight=0

L15 prim_mst: seen.add(u)
n=2; edges=[]; adj=[[], []]; heap=[]; seen={0}; total=0; u=0; weight=0

Why the solution is correct

After the root, each accepted entry is a minimum-weight edge crossing from the current tree to an unseen vertex: every such edge has been offered to the heap. The MST cut property makes this extension safe. If the heap empties before all vertices are included, no edge connects the remaining components.

Time and space

O(V+E log(E+2)) time and O(V+E) space for this lazy edge-heap implementation, including adjacency.

Common mistakes

Adding a stale edge's weight before checking seen double-counts vertices. Heap priorities are individual edge weights, unlike Dijkstra's accumulated path distances.

Reasoning card

Transfer exercise

Record parent edges to return the actual tree, then compare total cost with Kruskal on the same graph.