Minimum Spanning Tree Prim's Algorithm in DSA Python

Mental Model

Start from any node and keep adding the cheapest edge that connects a new node to the growing tree until all nodes are included.

Analogy: Imagine building a road network connecting all cities with the least total cost by always choosing the shortest new road from the cities already connected.

Graph:
  (0)
 /   \
(1)---(2)
Edges with weights:
0-1: 4, 0-2: 3, 1-2: 5
Start at node 0 ↑

Dry Run Walkthrough

Input: Graph with nodes 0,1,2 and edges: 0-1(4), 0-2(3), 1-2(5). Start from node 0.

Goal: Build a tree connecting all nodes with minimum total edge weight.

Step 1: Start from node 0, add it to MST set

MST nodes: [0]
Edges considered: 0-1(4), 0-2(3)

Why: We begin from node 0 to start building the MST.

Step 2: Choose edge with smallest weight from MST to non-MST nodes: 0-2(3)

MST nodes: [0, 2]
Edges considered: 0-1(4), 2-1(5)

Why: Adding node 2 via edge 0-2 keeps total cost minimal.

Step 3: Choose next smallest edge connecting MST to non-MST: 0-1(4)

MST nodes: [0, 2, 1]
All nodes included, MST complete

Why: Add node 1 with edge 0-1 to connect all nodes with minimal cost.

Result:

MST edges: 0-2(3) -> 0-1(4)
Total cost: 7

Annotated Code

DSA Python

import heapq

class Graph:
    def __init__(self, vertices):
        self.V = vertices
        self.graph = [[] for _ in range(vertices)]

    def add_edge(self, u, v, w):
        self.graph[u].append((v, w))
        self.graph[v].append((u, w))

    def prim_mst(self):
        mst_set = [False] * self.V  # Track nodes included in MST
        edge_heap = []  # Min-heap for edges (weight, node)
        result = []  # Store MST edges

        # Start from node 0
        mst_set[0] = True
        for v, w in self.graph[0]:
            heapq.heappush(edge_heap, (w, 0, v))  # Push edges from node 0

        while edge_heap and len(result) < self.V - 1:
            w, u, v = heapq.heappop(edge_heap)  # Get smallest edge
            if not mst_set[v]:
                mst_set[v] = True  # Add node to MST
                result.append((u, v, w))  # Add edge to MST
                for to, weight in self.graph[v]:
                    if not mst_set[to]:
                        heapq.heappush(edge_heap, (weight, v, to))  # Add new edges
        return result

# Driver code
if __name__ == '__main__':
    g = Graph(3)
    g.add_edge(0, 1, 4)
    g.add_edge(0, 2, 3)
    g.add_edge(1, 2, 5)

    mst = g.prim_mst()
    total_cost = 0
    for u, v, w in mst:
        print(f"Edge {u}-{v} with weight {w}")
        total_cost += w
    print(f"Total cost: {total_cost}")

mst_set[0] = True

Mark starting node 0 as included in MST

for v, w in self.graph[0]: heapq.heappush(edge_heap, (w, 0, v))

Add all edges from node 0 to the heap to consider

w, u, v = heapq.heappop(edge_heap)

Select edge with smallest weight connecting MST to new node

if not mst_set[v]: mst_set[v] = True result.append((u, v, w))

Add new node and edge to MST if node not already included

for to, weight in self.graph[v]: if not mst_set[to]: heapq.heappush(edge_heap, (weight, v, to))

Add edges from newly included node to heap for future selection

Output

Edge 0-2 with weight 3

Edge 0-1 with weight 4

Total cost: 7

Complexity Analysis

Time: O(E log V) because each edge is pushed and popped from the heap at most once, and heap operations take log V time

Space: O(V + E) to store the graph and the heap data structures

vs Alternative: Compared to Kruskal's algorithm which sorts edges first, Prim's uses a heap and adjacency list, often more efficient for dense graphs

Edge Cases

Graph with a single node and no edges

MST is empty since no edges exist

DSA Python

while edge_heap and len(result) < self.V - 1:

Graph with disconnected nodes

Algorithm only builds MST for connected component containing start node; others remain disconnected

DSA Python

while edge_heap and len(result) < self.V - 1:

Graph with multiple edges of same weight

Algorithm picks edges in order of heap; ties resolved by heap order but MST remains minimal

DSA Python

w, u, v = heapq.heappop(edge_heap)

When to Use This Pattern

When you need to connect all points with minimum total cost and the graph is weighted, use Prim's MST algorithm to grow the tree step-by-step by choosing the cheapest connecting edge.

Common Mistakes

Mistake: Not marking nodes as included in MST before adding their edges to the heap
Fix: Mark the node as included immediately after selecting it to avoid revisiting

Mistake: Adding edges to the heap without checking if the connected node is already in MST
Fix: Only add edges to nodes not yet included in MST to prevent cycles

Summary

Builds a minimum spanning tree by adding the cheapest edge connecting new nodes to the growing tree.

Use when you want to connect all nodes in a weighted graph with the least total edge cost.

Always pick the smallest edge from the current tree to a new node to keep the total cost minimal.

Practice

1. What is the main goal of Prim's algorithm in graph theory?

easy

A. To detect cycles in a graph

B. To find the shortest path between two vertices

C. To find the minimum spanning tree connecting all vertices with the least total edge weight

D. To sort the vertices in topological order

Solution

1. Step 1: Understand Prim's algorithm purpose

   Prim's algorithm builds a tree connecting all vertices with minimum total edge weight.

2. Step 2: Compare with other graph algorithms

   Shortest path, cycle detection, and topological sort serve different purposes than Prim's.

3. Final Answer:

   To find the minimum spanning tree connecting all vertices with the least total edge weight -> Option C

4. Quick Check:

   Prim's algorithm = Minimum Spanning Tree [OK]

Hint: Prim's finds minimum spanning tree, not shortest path [OK]

Common Mistakes:

• Confusing Prim's with Dijkstra's shortest path
• Thinking it detects cycles
• Mixing it with topological sorting

2. Which of the following Python data structures is most suitable to efficiently select the next edge with the smallest weight in Prim's algorithm?

easy

A. List

B. Priority Queue (heapq)

C. Dictionary

D. Queue

Solution

1. Step 1: Identify the need for selecting minimum edge

   Prim's algorithm repeatedly selects the edge with the smallest weight.

2. Step 2: Choose data structure for efficient minimum extraction

   Priority Queue (heapq) allows efficient extraction of the smallest element.

3. Final Answer:

   Priority Queue (heapq) -> Option B

4. Quick Check:

   Minimum edge selection = Priority Queue [OK]

Hint: Use heapq for fast minimum edge extraction [OK]

Common Mistakes:

• Using list and scanning all edges each time
• Using queue which is FIFO, not minimum-based
• Using dictionary which doesn't order by weight

3. What is the output of the following Python code implementing Prim's algorithm on the given graph?

graph = {
  0: [(1, 2), (3, 6)],
  1: [(0, 2), (2, 3), (3, 8), (4, 5)],
  2: [(1, 3), (4, 7)],
  3: [(0, 6), (1, 8), (4, 9)],
  4: [(1, 5), (2, 7), (3, 9)]
}

import heapq

def prim(graph, start=0):
    mst = []
    visited = set([start])
    edges = [(weight, start, to) for to, weight in graph[start]]
    heapq.heapify(edges)

    while edges:
        weight, frm, to = heapq.heappop(edges)
        if to not in visited:
            visited.add(to)
            mst.append((frm, to, weight))
            for to_next, w in graph[to]:
                if to_next not in visited:
                    heapq.heappush(edges, (w, to, to_next))
    return mst

print(prim(graph))

medium

A. [(0, 1, 2), (1, 2, 3), (1, 4, 5), (0, 3, 6)]

B. [(0, 3, 6), (3, 4, 9), (1, 4, 5), (1, 2, 3)]

C. [(0, 1, 2), (0, 3, 6), (1, 4, 5), (2, 4, 7)]

D. [(1, 2, 3), (2, 4, 7), (0, 3, 6), (1, 4, 5)]

Solution

1. Step 1: Start from vertex 0 and add edges

   Edges from 0: (0-1,2), (0-3,6). Choose smallest: (0-1,2).

2. Step 2: Add edges from vertex 1 and pick next smallest

   From 1: (1-2,3), (1-3,8), (1-4,5). Next smallest edge overall is (1-2,3).

3. Step 3: Continue adding edges and selecting minimum

   Next edges: (1-4,5), (0-3,6). Choose (1-4,5) then (0-3,6).

4. Final Answer:

   [(0, 1, 2), (1, 2, 3), (1, 4, 5), (0, 3, 6)] -> Option A

5. Quick Check:

   Edges chosen in ascending weight order form MST [OK]

Hint: Pick edges by smallest weight connecting new vertices [OK]

Common Mistakes:

• Choosing edges with higher weight first
• Adding edges that create cycles
• Not updating visited vertices correctly

4. The following code snippet for Prim's algorithm has a bug. What is the error?

def prim_bug(graph, start=0):
    mst = []
    visited = set()
    edges = [(0, start, start)]
    while edges:
        weight, frm, to = edges.pop()
        if to not in visited:
            visited.add(to)
            if frm != to:
                mst.append((frm, to, weight))
            for to_next, w in graph[to]:
                if to_next not in visited:
                    edges.append((w, to, to_next))
    return mst

medium

A. The start vertex is never added to visited

B. Visited set is not initialized properly

C. Edges are not added to the list

D. Using list pop() instead of priority queue causes wrong edge selection order

Solution

1. Step 1: Check how edges are selected

   Edges are popped from a list using pop(), which removes the last element, not the smallest weight.

2. Step 2: Understand Prim's requirement for minimum edge selection

   Prim's requires selecting the edge with the smallest weight each time, which needs a priority queue (heap).

3. Final Answer:

   Using list pop() instead of priority queue causes wrong edge selection order -> Option D

4. Quick Check:

   Edge selection order must be by minimum weight [OK]

Hint: Use heapq, not list pop(), to pick minimum edge [OK]

Common Mistakes:

• Using list pop() instead of heapq.heappop()
• Not initializing visited with start vertex
• Appending edges incorrectly

5. Given a weighted undirected graph with vertices 0 to 5 and edges:

0-1: 4, 0-2: 3, 1-2: 1, 1-3: 2, 2-3: 4, 3-4: 2, 4-5: 6

Using Prim's algorithm starting at vertex 0, what is the total weight of the minimum spanning tree?

hard

A. 9

B. 14

C. 10

D. 12

Solution

1. Step 1: Start from vertex 0 and pick smallest edges

   Edges from 0: (0-2,3), (0-1,4). Choose (0-2,3).

2. Step 2: Add edges from visited vertices and pick smallest

   Visited: {0,2}. Edges: (2-1,1), (0-1,4), (2-3,4). Choose (2-1,1).

3. Step 3: Continue adding edges

   Visited: {0,1,2}. Edges: (1-3,2), (0-1,4), (2-3,4). Choose (1-3,2).

4. Step 4: Add edges from 3

   Visited: {0,1,2,3}. Edges: (3-4,2), (2-3,4), (0-1,4). Choose (3-4,2).

5. Step 5: Add edges from 4

   Visited: {0,1,2,3,4}. Edges: (4-5,6), (2-3,4), (0-1,4). Choose (4-5,6).

6. Step 6: Calculate total weight

   Edges chosen: 3 + 1 + 2 + 2 + 6 = 14.

7. Final Answer:

   14 -> Option A

8. Quick Check:

   Sum of MST edges = 14 [OK]

Hint: Add edges by smallest weight connecting new vertices until all included [OK]

Common Mistakes:

• Including edges that form cycles
• Skipping the smallest connecting edge
• Miscounting total weight