Prim's algorithm versus Kruskal's in Python

Why it matters

A company wants to connect five offices with network cable. Each pair of offices can be joined by a cable with a known cost, but the company does not need every pair connected. It needs enough cables that every office can reach every other one, with the smallest total cost. A road builder has the same problem with roads between towns, and a chip designer has it with wires between parts of a circuit.

The answer is a minimum spanning tree, usually shortened to MST. A spanning tree is a set of edges that connects every vertex without a cycle. A cycle is a path that starts and ends at the same vertex. The minimum spanning tree is the spanning tree with the smallest total weight. The weight of an edge is its cost, such as the length of a cable.

Two algorithms find it, and both are greedy. A greedy algorithm makes the locally best choice at each step and never goes back. Prim's algorithm grows one tree from a starting vertex, one cheapest edge at a time. Kruskal's algorithm sorts all the edges and takes them in order, skipping any edge that would close a cycle. Both give a minimum spanning tree, and they do it in different ways.

The module already has a lesson on union-find with an early version of Kruskal's and Prim's algorithms, 05-union-find-and-minimum-spanning-trees. This lesson puts the two side by side, checks them against each other on random graphs, and adds Prim's algorithm for dense graphs. The deeper version of union-find, with its speed tricks, is the next lesson, so this lesson uses the simplest form.

You should know how to read numbers from input, how to write a function, and how a list works. The Python module heapq is a standard library module that keeps a list in heap order. If you have used Java, heapq does the job of PriorityQueue. Every program in this lesson is a complete Python file. The Input box text is shown before each program that reads input.

Words used in this lesson

A vertex is a point in the graph, such as an office. An edge joins two vertices and has a weight, which is its cost.

A tree is a connected graph with no cycles. A tree on n vertices always has exactly n - 1 edges. This fact is used in the code to know when the job is done.

A cut is a way of splitting the vertices into two groups, inside and outside. An edge crosses the cut if one end is inside and the other end is outside. Both algorithms depend on one simple fact about cuts, explained in the next section.

Dense and sparse describe how many edges a graph has. A sparse graph has about as many edges as vertices. A dense graph has close to the largest possible number of edges, which is about V * (V - 1) / 2 for V vertices. Office networks are often sparse, because each office is cabled to a few neighbours. A graph where every office is connected to every other one is dense.

The sample graph

Use this graph of five offices, numbered 0 to 4, for the whole lesson. Each line is one cable, written as u v w, which means it joins vertex u and vertex v with cost w. The cables work both ways.

• 0 to 1, cost 2
• 0 to 3, cost 6
• 1 to 2, cost 3
• 1 to 3, cost 8
• 1 to 4, cost 5
• 2 to 4, cost 7
• 3 to 4, cost 9

The tree we want has four cables, because there are five offices. A good tree uses the cheap cables and avoids the expensive ones. We will check that both algorithms find a tree of total cost 16.

The key fact behind both algorithms

Pick any cut, which is any split of the offices into two groups. Look at all the cables that cross the cut. The cheapest of those crossing cables is always safe to put into a minimum spanning tree.

Here is why, in plain words. Suppose the cheapest crossing cable is not in the best tree. The best tree must connect the two groups somehow, so it uses some other crossing cable, which costs at least as much. Swap that more expensive cable for the cheap one. The tree still connects everything, and it costs no more. So a cheapest crossing cable can always be used.

Prim's algorithm uses the cut "the offices already in the tree, and the offices not yet in the tree". Kruskal's algorithm uses cuts that come from the groups it has built so far. Both algorithms repeatedly take the cheapest edge that crosses a cut, which is why they are correct.

Prim's algorithm: grow one tree

Prim's algorithm starts with one office in the tree. At each step it looks at every cable that leaves the tree, that is, every cable from an office in the tree to an office outside it. It takes the cheapest one and adds its outside office to the tree. It repeats until every office is in the tree.

A heap, which always gives back its smallest item first, makes the "cheapest cable that leaves the tree" fast to find. Put every cable from the tree into a heap, ordered by cost. In Python this is a list managed by heapq.heappush and heapq.heappop. The smallest item is at the front, and heappop removes it. If its far end is already in the tree, the cable is no longer useful, so skip it. Otherwise add its far end to the tree, and push the new cables that leave that office.

The programs store each cable as an Edge named tuple with the weight first. A named tuple is a tuple whose fields have names. Python compares tuples one field at a time, from the first field. Putting w first means that edges compare by cost, so the heap and sorted need no extra rule.

Stale entries are normal. A cable can be pushed into the heap while its far end is outside the tree. Later, a cheaper path brings that far end into the tree. The old cable is still in the heap, and it now points inside the tree. Popping it and skipping it is the fix, and it is the same idea as the stale entries in the Dijkstra lesson.

A worked trace of Prim's algorithm

Start with the tree {0}. The cables leaving 0 go to 1 (cost 2) and to 3 (cost 6). The heap holds 0 to 1 (2) and 0 to 3 (6).

Step 1. Pop the cable 0 to 1, cost 2. Office 1 is outside the tree, so add it. The tree is {0, 1}. The cables leaving office 1 are 1 to 2 (cost 3), 1 to 3 (cost 8) and 1 to 4 (cost 5). The cable back to 0 is inside the tree, so it is not pushed. The heap now holds 0 to 3 (6), 1 to 2 (3), 1 to 4 (5) and 1 to 3 (8).

Step 2. Pop 1 to 2, cost 3. Office 2 is outside, so add it. The tree is {0, 1, 2}. The cable 2 to 4 (cost 7) leaves the tree and is pushed. The cable back to 1 is inside, so it is not pushed. The heap holds 1 to 4 (5), 0 to 3 (6), 2 to 4 (7) and 1 to 3 (8).

Step 3. Pop 1 to 4, cost 5. Office 4 is outside, so add it. The tree is {0, 1, 2, 4}. The cable 4 to 3 (cost 9) is pushed. The cable 4 to 2 is inside, so the entry for it is not needed. The heap holds 0 to 3 (6), 1 to 3 (8), 2 to 4 (7) and 4 to 3 (9).

Step 4. Pop 0 to 3, cost 6. Office 3 is outside, so add it. The tree is {0, 1, 2, 3, 4}, which has five offices and four cables. The total cost is 2 + 3 + 5 + 6 = 16.

The heap still holds three entries: 1 to 3 (8), 2 to 4 (7) and 4 to 3 (9). Each one points inside the tree now, so they are stale. The program stops when it has four cables, and it never pops them.

Kruskal's algorithm: sort the cables and join groups

Kruskal's algorithm takes a different view. It does not grow one tree. It starts with every office alone, as its own group. Then it goes through all the cables from cheapest to most expensive. For each cable, it asks whether the two ends are in different groups. If they are, the cable joins the two groups, and it is added to the tree. If they are already in the same group, the cable would close a cycle, so it is skipped.

Two tools make this work. The first is a sort, which puts the cables in order of cost. The second is union-find, a structure that answers "which group is this office in?" and joins two groups in one step. Union-find is explained in the earlier lesson. Here we use the simplest version: each office points to a parent, the root of the tree of parents is the group's name, and find walks up the parents to reach the root.

A worked trace of Kruskal's algorithm

Sort the cables by cost. The order is 0 to 1 (2), 1 to 2 (3), 1 to 4 (5), 0 to 3 (6), 2 to 4 (7), 1 to 3 (8), 3 to 4 (9).

Start with five groups: {0}, {1}, {2}, {3}, {4}.

Cable 0 to 1, cost 2. The ends are in different groups. Join them: {0, 1}. Add the cable.

Cable 1 to 2, cost 3. Office 2 is in its own group, and office 1 is in {0, 1}. Join: {0, 1, 2}. Add the cable.

Cable 1 to 4, cost 5. Office 4 is alone. Join: {0, 1, 2, 4}. Add the cable.

Cable 0 to 3, cost 6. Office 3 is alone. Join: {0, 1, 2, 3, 4}. Add the cable. There are now four cables, which is n - 1, so the job is done. The program stops here.

The cables 2 to 4 (cost 7), 1 to 3 (cost 8) and 3 to 4 (cost 9) are never examined, because the loop stops once four cables are taken. Kruskal skips no cable on this graph. The total is again 2 + 3 + 5 + 6 = 16.

Program 1: both algorithms on the sample graph

This program reads the graph and runs both algorithms on it. It prints the total cost and the cables each one chose.

Type this into the Input box:

5 7
0 1 2
0 3 6
1 2 3
1 3 8
1 4 5
2 4 7
3 4 9

The first line says there are 5 vertices and 7 edges. Each of the next seven lines is one edge.

import heapq
import sys
from typing import NamedTuple
 
 
class Edge(NamedTuple):
    w: int  # weight first, so edges compare by cost
    u: int
    v: int
 
 
def build_adjacency(n, edges):
    adj = [[] for _ in range(n)]
    for e in edges:
        adj[e.u].append(e)
        adj[e.v].append(Edge(e.w, e.v, e.u))
    return adj
 
 
def prim(n, edges, start=0):
    adj = build_adjacency(n, edges)
    in_tree = [False] * n
    in_tree[start] = True
    heap = list(adj[start])
    heapq.heapify(heap)
    chosen = []
    total = 0
    while heap and len(chosen) < n - 1:
        e = heapq.heappop(heap)
        if in_tree[e.v]:
            continue  # both ends are already in the tree
        in_tree[e.v] = True
        chosen.append(e)
        total += e.w
        for f in adj[e.v]:
            if not in_tree[f.v]:
                heapq.heappush(heap, f)
    if len(chosen) == n - 1:
        return total, chosen
    return None, chosen
 
 
def find(parent, x):
    while parent[x] != x:
        x = parent[x]
    return x
 
 
def kruskal(n, edges):
    parent = list(range(n))
    chosen = []
    total = 0
    for e in sorted(edges):
        a = find(parent, e.u)
        b = find(parent, e.v)
        if a != b:
            parent[a] = b
            chosen.append(e)
            total += e.w
            if len(chosen) == n - 1:
                break
    if len(chosen) == n - 1:
        return total, chosen
    return None, chosen
 
 
def main():
    nums = list(map(int, sys.stdin.read().split()))
    n, m = nums[0], nums[1]
    edges = []
    for i in range(m):
        u, v, w = nums[2 + 3 * i: 5 + 3 * i]
        edges.append(Edge(w, u, v))
 
    prim_total, prim_edges = prim(n, edges)
    print("Prim total:", prim_total, "edges:", prim_edges)
 
    kruskal_total, kruskal_edges = kruskal(n, edges)
    print("Kruskal total:", kruskal_total, "edges:", kruskal_edges)
 
 
if __name__ == "__main__":
    main()

Expected output:

Prim total: 16 edges: [Edge(w=2, u=0, v=1), Edge(w=3, u=1, v=2), Edge(w=5, u=1, v=4), Edge(w=6, u=0, v=3)]
Kruskal total: 16 edges: [Edge(w=2, u=0, v=1), Edge(w=3, u=1, v=2), Edge(w=5, u=1, v=4), Edge(w=6, u=0, v=3)]

Both totals are 16, as the traces predicted. The edges chosen are the same here, in the same order. In general the two algorithms can choose different edges when two cables have the same cost, but both always reach the same total. The Edge(...) text in the output is how Python prints a named tuple. The field order is w, u, v, so the weight comes first.

Look at the line len(chosen) < n - 1 in prim and len(chosen) == n - 1 in kruskal. These are the stopping rules. A tree on n offices has n - 1 cables, so the program can stop as soon as it has that many. The None return is for a graph that is not connected, where no tree of that size exists.

The function prim also takes a start argument. It is 0 by default, and Exercise 1 changes it.

Dense and sparse graphs: which one to pick

The two algorithms have different costs, and the cost depends on how many edges the graph has. Let V be the number of vertices and E the number of edges. The notation O(...) shows how the number of steps grows as the graph grows, and log V is roughly how many steps one heap operation takes.

Prim with a heap does about E pushes and pops, and each one costs log V. Its time is O(E log V). Kruskal sorts all E edges first, which costs O(E log E), and then does almost constant work per edge with union-find. Its time is O(E log E).

For a sparse graph, where E is about V, both are fast. The sort in Kruskal is small, and Prim's heap is small too. Kruskal is simple to write and easy to check, so it is a good choice for sparse graphs.

For a dense graph, where E is about V * V / 2, the costs grow quickly. Kruskal must sort about V * V / 2 edges, which is O(V^2 log V). Prim with a heap does about the same amount of work, because log E is only about twice log V when E is close to V * V. So a heap does not help much on a dense graph.

Prim with an adjacency matrix avoids the heap entirely. An adjacency matrix is a two-dimensional list w[u][v] that holds the cost of the edge between u and v, or a sentinel value when there is no edge. A sentinel is a special value that means "not a real value". Prim with a matrix looks for the cheapest unused vertex with one linear scan, which takes O(V), and it does this V times. Its time is O(V^2), with no log factor. For a dense graph this is the best of the three, because the matrix is already the shape of a dense graph.

The table below gives the simple comparison. Use it as a guide, not as a rule that replaces measuring your own input.

Program 2: a dense matrix and a random check

The second program has three parts. It runs Prim with a matrix, which is the dense-graph method. It runs the other two methods on random connected graphs, and it checks that all three give the same total. This is a good way to test a new algorithm, because the three methods are independent of one another. If one is wrong, it will usually disagree with the other two.

This program does not read input. It creates its own graphs with a fixed random seed, so the output is the same on every run.

Expected output:

checked 300 graphs, mismatches 0
dense Prim on the sample graph: 16

The random graphs have between 2 and 26 vertices. Each one starts with a spanning tree, which guarantees that the graph is connected, and then gets a random number of extra edges. The matrix stores the cheapest cost between each pair, because a random graph can have two edges between the same pair of vertices. The min calls keep the cheaper one. Without that, the matrix would keep only the last edge read, which can be the more expensive one.

The no-edge value is math.inf, which is Python's floating-point infinity. Every real cost is smaller than infinity, so the scan never picks a missing edge while a real one is available. Python integers can be as large as you need, so there is no fixed largest integer to use as a sentinel, and math.inf is the usual choice in a matrix like this.

The line if not (a == b == c) is a chained comparison. It means a == b and b == c. Python lets you write the chain in one expression.

The fixed seed 2026 makes the same 300 graphs appear on every run. A test with a fixed seed is easy to repeat, and a bug can be reproduced by running the same seed again.

Common mistakes

Mistake 1: forgetting the stale-entry check in Prim. Without the check, the program adds a cable whose far end is already in the tree. That closes a cycle, and the total becomes too small or too large. Here is a graph where the bug shows. It has four offices and five cables, and the cable from 0 to 2 costs 5, but offices 0 and 2 are already joined through 1 at a cost of 3.

Type this into the Input box:

4 5
0 1 1
0 2 5
1 2 2
2 3 9
0 3 20

The correct program, from Program 1, prints:

Prim total: 12 edges: [Edge(w=1, u=0, v=1), Edge(w=2, u=1, v=2), Edge(w=9, u=2, v=3)]

Delete these two lines from prim in Program 1:

        if in_tree[e.v]:
            continue  # both ends are already in the tree

Then the output is:

Prim total: 8 edges: [Edge(w=1, u=0, v=1), Edge(w=2, u=1, v=2), Edge(w=5, u=0, v=2)]

The total 8 is wrong. The chosen cables include 0 to 2, which closes a cycle with 0 to 1 and 1 to 2, and the cable 2 to 3 is missing. A wrong answer like this is harder to notice than a crash, so test against a second method, as Program 2 does.

Mistake 2: forgetting to sort in Kruskal. Kruskal's method depends on taking the cheapest cable first. If sorted(edges) is replaced by edges, the cables are taken in the order they were typed. On the same four-office graph, the program takes 0 to 1 (cost 1), then 0 to 2 (cost 5), then skips 1 to 2 because it closes a cycle, then takes 2 to 3 (cost 9). The total is 15, not 12. The output is:

Kruskal total: 15 edges: [Edge(w=1, u=0, v=1), Edge(w=5, u=0, v=2), Edge(w=9, u=2, v=3)]

The sort is one call, and it is the whole idea of the algorithm. Test that it runs before the union-find part.

Mistake 3: putting items in a heap that cannot be ordered. A heap compares its items, so each item must support <. A plain class does not, because Python does not know which of two objects is smaller. The class below has no ordering rule, and pushing two of its objects into a heap fails.

import heapq
from dataclasses import dataclass
 
 
@dataclass
class Edge:
    u: int
    v: int
    w: int
 
 
heap = []
heapq.heappush(heap, Edge(0, 1, 2))
heapq.heappush(heap, Edge(0, 2, 5))

Running it stops with an error. The last line of the message says what went wrong:

TypeError: '<' not supported between instances of 'Edge' and 'Edge'

There are two fixes. Use a named tuple with the weight first, as Program 1 does, which is the simplest for an edge. Or add order=True to the decorator, as in @dataclass(order=True). The fields then compare in the order they are declared, so w must come first in the class.

Mistake 4: treating a missing edge as cost 0 in the matrix. In the matrix, the cost of a missing edge must be a value that no real edge can have. Here it is math.inf. If the code uses 0 for "no edge", Prim picks the missing edge first, because 0 is the cheapest possible cost, and the total becomes wrong. On a 3-vertex graph with edges 0 to 1 (cost 4) and 1 to 2 (cost 7), the correct total is 11. With 0 as the no-edge value, the output is 4:

with 0 as no-edge: 4
with inf as no-edge: 11

Use a large sentinel, and check the key before adding it to the total, as the return None in prim_dense does.

Mistake 5: not checking connectivity. Both programs return None when the graph is not connected. If the code ignores the None, it uses a value that means nothing. A graph with two separate parts has no spanning tree, so the honest answer is "no tree exists", not a number.

Practice

Try each exercise by hand first, then check with the programs.

Exercise 1 (easy). Run Prim on the sample graph, but start from vertex 2 instead of vertex 0. Use prim(5, edges, start=2). Which four cables are chosen, and what is the total?

Show the answer

The cables are 2 to 1 (cost 3), 1 to 0 (cost 2), 1 to 4 (cost 5) and 0 to 3 (cost 6). The total is 16, the same as before. The program prints each cable in the order it joins the tree, so the first one appears as Edge(w=3, u=2, v=1). The starting vertex changes the order in which cables are added, but not the final total.

Exercise 2 (easy). For the sample graph, Kruskal's algorithm never examines three cables. Name them and say why.

Show the answer

The cables 2 to 4 (cost 7), 1 to 3 (cost 8) and 3 to 4 (cost 9) come after the fourth cable in sorted order. Kruskal takes 0 to 3 (cost 6) as the fourth cable, so the loop stops before it reaches them.

Exercise 3 (medium). Find the minimum spanning tree and its total for this graph with four vertices: 0 to 1 cost 1, 0 to 2 cost 10, 1 to 2 cost 2, 2 to 3 cost 3, and 1 to 3 cost 4.

Show the answer

The cables are 0 to 1 (1), 1 to 2 (2) and 2 to 3 (3). The total is 6. The cable 1 to 3 (cost 4) closes a cycle with 1 to 2 and 2 to 3, and the cable 0 to 2 (cost 10) is more expensive than the other route, so it is not needed.

Exercise 4 (medium). Explain why Prim with an adjacency matrix does not need a heap, and why it is a good choice for a dense graph.

Show the answer

A matrix gives direct access to the cost between any two vertices, so Prim can look for the cheapest unused vertex with one scan of V values. That scan costs O(V), and it runs V times, for O(V^2) in total. A dense graph has about V * V / 2 edges, so the matrix needs about that much space anyway, and the scan does not pay the extra log factor that a heap does. A sparse graph would waste most of the matrix on missing edges, which is why the heap version is used there.

Exercise 5 (hard). Program 1 starts Prim at vertex 0. What does prim return for a graph with one vertex and no edges, where n is 1?

Show the answer

The loop condition len(chosen) < n - 1 is 0 < 0, which is false, so the loop never runs. The return check len(chosen) == n - 1 is 0 == 0, which is true, so the function returns the pair (0, []). A single vertex is already a tree with no edges and cost 0, so the answer is correct.

Summary

• A minimum spanning tree connects every vertex with the smallest total edge cost and no cycle.
• Prim's algorithm grows one tree with heapq, and it skips stale entries whose far end is already in the tree.
• Kruskal's algorithm sorts all edges and takes each one that joins two different groups, using union-find to check the groups.
• Both algorithms rely on the rule that the cheapest edge crossing a cut is safe to add.
• On sparse graphs both are fast, and on dense graphs Prim with an adjacency matrix avoids the sort and the heap.
• Test a new implementation against a second method on random graphs, as Program 2 does.