DSA course menu
Minimum Spanning Trees: Kruskal's and Prim's Algorithms in Python
Learn minimum spanning trees: connecting every vertex of a weighted graph at the lowest total cost, Kruskal's algorithm with union-find, Prim's algorithm with a heap, and real uses like network design.
A spanning tree of a connected, undirected graph is a set of edges that connects all the vertices with no cycles. It always has exactly V − 1 edges. A minimum spanning tree (MST) is the spanning tree with the smallest total edge weight.
Think of laying cable between towns: every town must be connected, directly or through others, and you want the least total cable. That is an MST problem. It also appears in network design, clustering, and approximating harder problems like the travelling salesperson.
Both classic MST algorithms are greedy: they repeatedly take the cheapest safe edge. (For why greedy works here, see Greedy Algorithms.)
The example graph used below (towns A to E, costs in km):
A ─4─ B
│ \ │ \
1 3 2 5
│ \│ \
C ──6── D ─7─ EKruskal's algorithm#
- Sort all edges from cheapest to most expensive.
- Go through them in order. Take an edge if it connects two different groups (does not create a cycle); skip it otherwise.
- Stop after V − 1 edges.
"Are these already connected?" is exactly what union-find answers:
def kruskal(vertices, edges):
parent = {v: v for v in vertices}
def find(x):
while parent[x] != x:
parent[x] = parent[parent[x]]
x = parent[x]
return x
tree, total = [], 0
for w, u, v in sorted(edges):
ru, rv = find(u), find(v)
if ru != rv: # different groups: safe to add
parent[rv] = ru
tree.append((u, v, w))
total += w
if len(tree) == len(vertices) - 1:
break
return tree, total
vertices = ["A", "B", "C", "D", "E"]
edges = [(4, "A", "B"), (1, "A", "C"), (3, "A", "D"), (2, "B", "D"),
(5, "B", "E"), (6, "C", "D"), (7, "D", "E")]
tree, total = kruskal(vertices, edges)
print(tree)
print("total:", total)[('A', 'C', 1), ('B', 'D', 2), ('A', 'D', 3), ('B', 'E', 5)]
total: 11The edges A–B (4), C–D (6) and D–E (7) are skipped or never needed: each would either close a cycle or come after the tree is complete. Sorting dominates the cost: O(E log E).
Prim's algorithm#
Prim's grows one tree from a starting vertex. At each step it adds the cheapest edge that leaves the tree, using a heap of candidate edges:
import heapq
def prim(graph, start):
in_tree = {start}
heap = [(w, start, v) for v, w in graph[start]]
heapq.heapify(heap)
tree, total = [], 0
while heap and len(in_tree) < len(graph):
w, u, v = heapq.heappop(heap)
if v in in_tree:
continue # both ends already in the tree
in_tree.add(v)
tree.append((u, v, w))
total += w
for nxt, w2 in graph[v]:
if nxt not in in_tree:
heapq.heappush(heap, (w2, v, nxt))
return tree, total
graph = {v: [] for v in "ABCDE"}
for w, u, v in [(4, "A", "B"), (1, "A", "C"), (3, "A", "D"), (2, "B", "D"),
(5, "B", "E"), (6, "C", "D"), (7, "D", "E")]:
graph[u].append((v, w))
graph[v].append((u, w))
tree, total = prim(graph, "A")
print(tree)
print("total:", total)[('A', 'C', 1), ('A', 'D', 3), ('D', 'B', 2), ('B', 'E', 5)]
total: 11Prim's picks the same total weight, 11, though it discovers the edges in a different order. Its cost is O(E log V) with a heap.
Kruskal or Prim?#
If several edges share the same weight, a graph can have more than one MST, but they all have the same total weight.
Common mistakes#
Practice#
Practice exercise
Connect all points: points on a grid can be joined with cables whose cost is the Manhattan distance between them. Write min_cable(points) that returns the minimum total cost to connect all points.
Show hint
Every pair of points is a possible edge. Build all edges with their Manhattan distances and run Kruskal.
Show solution
def min_cable(points):
n = len(points)
edges = []
for i in range(n):
for j in range(i + 1, n):
(x1, y1), (x2, y2) = points[i], points[j]
edges.append((abs(x1 - x2) + abs(y1 - y2), i, j))
parent = list(range(n))
def find(x):
while parent[x] != x:
parent[x] = parent[parent[x]]
x = parent[x]
return x
total, used = 0, 0
for w, i, j in sorted(edges):
ri, rj = find(i), find(j)
if ri != rj:
parent[rj] = ri
total += w
used += 1
if used == n - 1:
break
return total
print(min_cable([(0, 0), (2, 2), (3, 10), (5, 2), (7, 0)]))
print(min_cable([(0, 0)]))20
0With n points there are n(n−1)/2 edges, so this is O(n² log n). Prim's algorithm with a simple array (no heap) solves the dense case in O(n²).
Key takeaway#
Next: exploring every possibility systematically, with Backtracking.
Finished this lesson?