Algorithms · Lesson 14 of 17
Minimum Spanning Trees: Kruskal and Prim
Learn minimum spanning trees in Python with Kruskal's algorithm and union-find, Prim's algorithm with a heap, and when to use each one.
- Advanced
- 18 min read
- 4 objectives
Before this lessonLesson 13: Topological Sort
What you will learn
- Define a spanning tree and a minimum spanning tree
- Build a union-find structure with path compression
- Implement Kruskal's and Prim's algorithms
- Compare their complexity and pick the right one
Your Progress
0 of 17 lessons 0%
- Lessons0 / 17
- Completed0
- Est. time left~ 5 hours
Create a free account to keep your progress on every device.
Tip: pressing Next marks this lesson complete automatically.
Imagine stackcone is wiring up offices in six cities. You know the cost of a fibre link between each pair of cities, and you want every office connected (directly or indirectly) for the least total cost. You do not need a direct link between every pair; you need just enough links that nothing is isolated, and no redundant loops.
That is a minimum spanning tree (MST). A spanning tree of a connected, undirected graph is a set of edges that connects all V nodes with no cycles, which always means exactly V - 1 edges. The minimum one has the smallest total weight. MSTs show up in network design, clustering (cut the most expensive MST edges to get groups) and as building blocks for approximation algorithms.
Note the difference from shortest paths: Dijkstra minimises the distance from one source to each node; an MST minimises the total cost of connecting everyone. The edges chosen can be quite different.
The example network
# (cost, city_a, city_b), undirected
links = [
(7, "Pune", "Mumbai"), (5, "Pune", "Hyderabad"), (8, "Mumbai", "Surat"),
(9, "Mumbai", "Hyderabad"), (7, "Hyderabad", "Bengaluru"), (15, "Surat", "Delhi"),
(6, "Hyderabad", "Delhi"), (8, "Bengaluru", "Delhi"), (11, "Surat", "Bengaluru"),
]
cities = sorted({c for _, a, b in links for c in (a, b)})
print(len(cities), "cities,", len(links), "possible links")6 cities, 9 possible links
The key fact: the cut property
Both algorithms are greedy, and both rest on one fact. Split the nodes into any two groups. The cheapest edge crossing between the groups is always safe to include in some MST. (If an MST skipped it, swapping it in for a more expensive crossing edge would give a cheaper tree, a contradiction.) Kruskal and Prim are just two different ways of picking the cut.
Union-find: the tool Kruskal needs
Kruskal needs to answer "are these two nodes already connected?" very quickly. A union-find (disjoint set) structure keeps each group as a tree with a root. find walks to the root; union links two roots. Two tricks, path compression and union by size, make each operation effectively constant time (technically O(α(n)), where α grows so slowly it is below 5 for any real input).
class UnionFind:
def __init__(self, items):
self.parent = {x: x for x in items}
self.size = {x: 1 for x in items}
def find(self, x):
while self.parent[x] != x:
self.parent[x] = self.parent[self.parent[x]] # path compression
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 connected
if self.size[ra] < self.size[rb]:
ra, rb = rb, ra
self.parent[rb] = ra # attach smaller under larger
self.size[ra] += self.size[rb]
return True
uf = UnionFind("abcd")
print(uf.union("a", "b"), uf.union("c", "d"), uf.union("b", "a"))
print(uf.find("a") == uf.find("b"), uf.find("a") == uf.find("c"))True True False True False
Kruskal's algorithm
Sort all edges by cost. Walk through them from cheapest to most expensive and keep an edge if it joins two different components (it does not create a cycle). Stop after V - 1 edges.
# from earlier in this lesson
# (cost, city_a, city_b), undirected
links = [
(7, "Pune", "Mumbai"), (5, "Pune", "Hyderabad"), (8, "Mumbai", "Surat"),
(9, "Mumbai", "Hyderabad"), (7, "Hyderabad", "Bengaluru"), (15, "Surat", "Delhi"),
(6, "Hyderabad", "Delhi"), (8, "Bengaluru", "Delhi"), (11, "Surat", "Bengaluru"),
]
cities = sorted({c for _, a, b in links for c in (a, b)})
class UnionFind:
def __init__(self, items):
self.parent = {x: x for x in items}
self.size = {x: 1 for x in items}
def find(self, x):
while self.parent[x] != x:
self.parent[x] = self.parent[self.parent[x]] # path compression
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 connected
if self.size[ra] < self.size[rb]:
ra, rb = rb, ra
self.parent[rb] = ra # attach smaller under larger
self.size[ra] += self.size[rb]
return True
# new code
def kruskal(nodes, edges):
uf = UnionFind(nodes)
tree, total = [], 0
for cost, a, b in sorted(edges):
if uf.union(a, b):
tree.append((a, b, cost))
total += cost
if len(tree) == len(nodes) - 1:
break
return tree, total
tree, total = kruskal(cities, links)
for a, b, cost in tree:
print(f"{a:10} - {b:10} {cost}")
print("total:", total)Pune - Hyderabad 5 Hyderabad - Delhi 6 Hyderabad - Bengaluru 7 Pune - Mumbai 7 Mumbai - Surat 8 total: 33
Time O(E log E), dominated by sorting; the union-find work is nearly linear. Space O(V + E). Kruskal is simple and shines on sparse graphs or when edges already arrive as a list.
Prim's algorithm
Prim grows a single tree from a starting node. At each step, add the cheapest edge that connects the tree to a node outside it. A min-heap of candidate edges makes "cheapest" fast. It looks a lot like Dijkstra, except the heap key is the edge cost alone, not the total distance from the start.
# from earlier in this lesson
# (cost, city_a, city_b), undirected
links = [
(7, "Pune", "Mumbai"), (5, "Pune", "Hyderabad"), (8, "Mumbai", "Surat"),
(9, "Mumbai", "Hyderabad"), (7, "Hyderabad", "Bengaluru"), (15, "Surat", "Delhi"),
(6, "Hyderabad", "Delhi"), (8, "Bengaluru", "Delhi"), (11, "Surat", "Bengaluru"),
]
cities = sorted({c for _, a, b in links for c in (a, b)})
# new code
import heapq
from collections import defaultdict
def prim(edges, start):
adj = defaultdict(list)
for cost, a, b in edges:
adj[a].append((cost, b))
adj[b].append((cost, a))
in_tree = {start}
heap = [(cost, start, nxt) for cost, nxt in adj[start]]
heapq.heapify(heap)
tree, total = [], 0
while heap and len(in_tree) < len(adj):
cost, frm, to = heapq.heappop(heap)
if to in in_tree:
continue # would form a cycle
in_tree.add(to)
tree.append((frm, to, cost))
total += cost
for c, nxt in adj[to]:
if nxt not in in_tree:
heapq.heappush(heap, (c, to, nxt))
return tree, total
tree, total = prim(links, "Delhi")
print([f"{a}-{b}" for a, b, _ in tree])
print("total:", total)['Delhi-Hyderabad', 'Hyderabad-Pune', 'Hyderabad-Bengaluru', 'Pune-Mumbai', 'Mumbai-Surat'] total: 33
Same total cost of 33, with edges discovered in a different order. Time O(E log V) with a binary heap, space O(V + E). Prim suits dense graphs stored as adjacency lists, and it never needs to sort every edge up front.
Kruskal or Prim?
- Edges given as a list, sparse graph: Kruskal is shortest to write.
- Dense graph, adjacency list or matrix: Prim avoids sorting all edges.
- Need connectivity queries anyway ("are these in the same group?"): Kruskal, because you already have union-find.
- Both give the same total cost; in interviews either is accepted if you explain the cut property.
Recap
- An MST connects all V nodes with V - 1 edges at minimum total cost.
- The cut property says the cheapest edge across any split is always safe, which makes greedy correct.
- Kruskal sorts edges and uses union-find to skip cycle-forming ones: O(E log E).
- Prim grows one tree using a heap of crossing edges: O(E log V).
- Union-find with path compression and union by size is nearly O(1) per operation.
# Write your solution here
Finished reading? Mark this lesson complete to track your progress.