Minimum Spanning Tree Prim's Algorithm in DSA Python - Build from Scratch

Minimum Spanning Tree using Prim's Algorithm

📖 Scenario: You are working as a network engineer. You need to connect several computers with cables. Each cable has a cost. Your goal is to connect all computers with the least total cost, without any loops.

🎯 Goal: Build a program that finds the minimum total cost to connect all computers using Prim's algorithm on a weighted graph.

📋 What You'll Learn

Create a graph as an adjacency list with exact nodes and edges

Use a variable to track the number of nodes

Implement Prim's algorithm to find the minimum spanning tree

Print the total cost of the minimum spanning tree

💡 Why This Matters

🌍 Real World

Network engineers use minimum spanning trees to design cost-efficient networks connecting computers or routers.

💼 Career

Understanding MST algorithms is important for roles in network design, infrastructure planning, and optimization.

1

Create the graph as an adjacency list

Create a dictionary called graph representing the weighted graph with these exact edges:
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)]

DSA Python

# Create the graph dictionary with exact edges
# Your code here

Hint

Use a dictionary where keys are node numbers and values are lists of tuples (neighbor, weight).

2

Set the number of nodes

Create a variable called num_nodes and set it to the number of nodes in the graph (which is 5).

DSA Python

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)]
}
# Create num_nodes variable and set it to 5
# Your code here

Hint

Count the keys in the graph dictionary.

3

Implement Prim's algorithm to find MST cost

Write code to find the total cost of the minimum spanning tree using Prim's algorithm. Use these exact variable names:
- visited as a set to track visited nodes,
- edges as a list to store edges to explore,
- total_cost to sum the MST cost.
Start from node 0 and use a while loop with condition len(visited) < num_nodes.
Use heapq for the priority queue and import it at the top.

DSA Python

import heapq

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)]
}
num_nodes = 5

# Implement Prim's algorithm to calculate total_cost
# Your code here

Hint

Use a priority queue to always pick the smallest edge connecting to unvisited nodes.

4

Print the total cost of the MST

Write a print statement to display the value of total_cost.

DSA Python

import heapq

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)]
}
num_nodes = 5

visited = set()
edges = []
total_cost = 0

visited.add(0)
for to_node, weight in graph[0]:
    heapq.heappush(edges, (weight, to_node))

while len(visited) < num_nodes:
    weight, node = heapq.heappop(edges)
    if node not in visited:
        visited.add(node)
        total_cost += weight
        for to_node, w in graph[node]:
            if to_node not in visited:
                heapq.heappush(edges, (w, to_node))

# Print the total_cost
# Your code here

Hint

Use print(total_cost) to show the final MST cost.

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