Minimum Spanning Tree in Python: Kruskal's and Prim's Algorithms

By Engr. Henry Akpan

Connecting Everything as Cheaply as Possible

A state government wants to connect seven towns with fibre-optic cable. Every town must be reachable from every other, but laying cable is expensive, so the total length should be as small as possible. A minimum spanning tree (MST) solves exactly this problem. It is used to design electricity grids, water pipes, computer networks and road layouts, and inside algorithms for clustering and image processing.

This tutorial covers minimum spanning trees in Python: what spanning trees are, the cut property that makes greedy algorithms work, Kruskal's algorithm (using union-find), Prim's algorithm (using a heap), their running times, a worked example traced step by step, and applications such as clustering. It builds on the tutorials on graphs, greedy algorithms, heaps and union-find.

What Is a Minimum Spanning Tree?

Consider a connected, undirected, weighted graph with V vertices. A spanning tree is a subset of its edges that:

• connects all the vertices, and
• contains no cycles.

Every spanning tree has exactly V − 1 edges. Fewer edges cannot connect everything, and one more edge would always create a cycle. A minimum spanning tree is a spanning tree with the smallest possible total edge weight. If several edges have equal weights, there may be more than one MST, but they all have the same total weight.

An MST is not the same as a set of shortest paths. Dijkstra's algorithm finds the shortest route from one starting point. An MST minimises the total cost of connecting everyone, so the path between two particular towns in the MST may be long.

Our Example Network

Seven towns, A to G, with the possible cable routes between them and their costs (in millions of naira):

edges = [
    (7, "A", "B"), (5, "A", "D"),
    (8, "B", "C"), (9, "B", "D"), (7, "B", "E"),
    (5, "C", "E"),
    (15, "D", "E"), (6, "D", "F"),
    (8, "E", "F"), (9, "E", "G"),
    (11, "F", "G"),
]
vertices = sorted({v for _, a, b in edges for v in (a, b)})
print(vertices, len(edges), "edges")

Output:

['A', 'B', 'C', 'D', 'E', 'F', 'G'] 11 edges

With seven towns, any spanning tree will have 7 − 1 = 6 edges.

Why Greedy Works: the Cut Property

Split the vertices into any two groups (a cut). The cheapest edge crossing between the two groups is always part of some MST. Both Kruskal's and Prim's algorithms repeatedly add such a "safe" cheapest crossing edge, which is why making the locally cheapest choice gives a globally optimal answer here, unlike many problems where greedy choices fail.

Kruskal's Algorithm

1. Sort all the edges from cheapest to most expensive.
2. Go through them in order. Add an edge if it joins two vertices that are not yet connected, and skip it if it would create a cycle.
3. Stop when V − 1 edges have been added.

Checking "are these two already connected?" is exactly what union-find does in effectively constant time.

class DSU:
    def __init__(self, items):
        self.parent = {x: x for x in items}

    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
        self.parent[rb] = ra
        return True

def kruskal(vertices, edges):
    dsu = DSU(vertices)
    tree, total = [], 0
    for w, a, b in sorted(edges):
        if dsu.union(a, b):
            tree.append((a, b, w))
            total += w
            if len(tree) == len(vertices) - 1:
                break
        else:
            print(f"  skip {a}-{b} ({w}): would form a cycle")
    return tree, total

tree, total = kruskal(vertices, edges)
for a, b, w in tree:
    print(f"{a}-{b}: {w}")
print("total cost:", total)

Output:

  skip B-C (8): would form a cycle
  skip E-F (8): would form a cycle
  skip B-D (9): would form a cycle
A-D: 5
C-E: 5
D-F: 6
A-B: 7
B-E: 7
E-G: 9
total cost: 39

The skip messages appear first because they are printed inside the loop, while the finished tree is printed afterwards. Here is the actual order of decisions:

1. A–D (5): add.
2. C–E (5): add.
3. D–F (6): add.
4. A–B (7): add.
5. B–E (7): add. This joins the {A, B, D, F} group with the {C, E} group.
6. B–C (8): skip, because B and C are already connected through E.
7. E–F (8): skip, because E and F are already connected.
8. B–D (9): skip, because B and D are already connected through A.
9. E–G (9): add. That makes 6 edges, so the algorithm stops.

(Edges of equal weight are taken in alphabetical order, because Python sorts the tuples by weight first and then by the vertex names.)

Running time: sorting the E edges takes O(E log E), and the union-find work is almost linear, so Kruskal's algorithm is O(E log E).

Prim's Algorithm

Prim's algorithm grows a single tree outwards from a starting vertex, like a spreading network:

1. Start the tree with any vertex.
2. Repeatedly add the cheapest edge that connects a vertex in the tree to a vertex outside it.
3. Stop when every vertex is in the tree.

A min-heap (heapq) keeps the candidate edges ordered, so the cheapest can always be taken quickly.

import heapq
from collections import defaultdict

def prim(edges, start):
    graph = defaultdict(list)
    for w, a, b in edges:
        graph[a].append((w, b))
        graph[b].append((w, a))

    in_tree = {start}
    heap = [(w, start, v) for w, v in graph[start]]
    heapq.heapify(heap)
    tree, total = [], 0
    while heap and len(in_tree) < len(graph):
        w, a, b = heapq.heappop(heap)
        if b in in_tree:
            continue                      # both ends already in the tree
        in_tree.add(b)
        tree.append((a, b, w))
        total += w
        for w2, nxt in graph[b]:
            if nxt not in in_tree:
                heapq.heappush(heap, (w2, b, nxt))
    return tree, total

tree, total = prim(edges, "A")
print(" -> ".join(f"{a}-{b}({w})" for a, b, w in tree))
print("total cost:", total)

Output:

A-D(5) -> D-F(6) -> A-B(7) -> B-E(7) -> E-C(5) -> E-G(9)
total cost: 39

Prim's algorithm picks the same six edges as Kruskal's, in a different order, and of course the same total of 39. Notice how the tree grows outwards from A: it reaches C only after the tree has reached E.

Running time: with a binary heap, Prim's algorithm is O(E log V).

Kruskal or Prim?

• Kruskal's works on a plain list of edges and is simple to write with union-find. It suits sparse graphs, with relatively few edges. It also works on disconnected graphs, giving a minimum spanning forest.
• Prim's works from an adjacency list and suits dense graphs, with many edges per vertex. It grows one connected tree.
• Both give a correct MST, and for most practical inputs either is fast enough.

Application: Clustering

A neat use of Kruskal's algorithm is clustering. To split points into k groups, run Kruskal's but stop when only k separate groups remain. Equivalently, build the MST and remove its k − 1 most expensive edges. The groups left are as far apart from each other as possible.

def clusters(vertices, edges, k):
    dsu = DSU(vertices)
    groups = len(vertices)
    for w, a, b in sorted(edges):
        if groups == k:
            break
        if dsu.union(a, b):
            groups -= 1
    result = defaultdict(list)
    for v in vertices:
        result[dsu.find(v)].append(v)
    return sorted(result.values())

print(clusters(vertices, edges, 2))
print(clusters(vertices, edges, 3))

Output:

[['A', 'B', 'C', 'D', 'E', 'F'], ['G']]
[['A', 'B', 'D', 'F'], ['C', 'E'], ['G']]

Town G is the most isolated (its cheapest link costs 9), so it is the first to form its own cluster.

Common Mistakes with Minimum Spanning Trees

• Confusing an MST with shortest paths. An MST minimises the total cost of connecting all the vertices, not the distance between two of them.
• Forgetting the cycle check in Kruskal's algorithm, which then adds edges that close loops.
• Running an MST algorithm on a directed graph. MSTs are defined for undirected graphs. The directed version (arborescence) needs a different algorithm.
• Not skipping stale heap entries in Prim's algorithm, the if b in in_tree: continue line, which adds the same vertex twice.
• Expecting a spanning tree of a disconnected graph. None exists. Kruskal's then gives a forest with fewer than V − 1 edges.

Minimum Spanning Trees: Practice Exercises

1. Modify kruskal so that it reports whether the graph is connected (the tree has V − 1 edges).
2. Find the MST weight of the graph with edges (1, P, Q), (3, P, R), (2, Q, R), (4, Q, S) and (5, R, S).
3. A maximum spanning tree has the largest possible total weight. How can you find one with Kruskal's algorithm?

Sample solutions:

def mst_info(vertices, edges):
    dsu = DSU(vertices)
    used, total = 0, 0
    for w, a, b in sorted(edges):
        if dsu.union(a, b):
            used += 1
            total += w
    return total, used == len(vertices) - 1

small = [(1, "P", "Q"), (3, "P", "R"), (2, "Q", "R"), (4, "Q", "S"), (5, "R", "S")]
print(mst_info(["P", "Q", "R", "S"], small))
print(mst_info(["P", "Q", "R", "S", "T"], small))       # T is unreachable

def max_spanning_weight(vertices, edges):
    negated = [(-w, a, b) for w, a, b in edges]
    total, _ = mst_info(vertices, negated)
    return -total

print(max_spanning_weight(["P", "Q", "R", "S"], small))

Output:

(7, True)
(7, False)
12

For question 3, sort the edges from most to least expensive. Negating the weights, as above, does exactly that.

Conclusion

A minimum spanning tree connects all V vertices of a weighted, undirected graph with V − 1 edges and the least total weight. The cut property guarantees that greedy choices work. Kruskal's algorithm sorts the edges and adds each one that does not form a cycle, checked with union-find, in O(E log E). Prim's algorithm grows one tree from a start vertex, always taking the cheapest edge leaving the tree, using a heap, in O(E log V). MSTs design cheap networks of cables, pipes and roads, and stopping Kruskal's early gives a simple clustering method.