Minimum Spanning Tree Prim's Algorithm in DSA Go

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)(3)

Start at node 0, pick cheapest edges to connect all nodes without cycles.

Dry Run Walkthrough

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

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

Step 1: Start at node 0, add node 0 to MST set

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

Why: We begin from node 0 and look at edges from it to find the cheapest connection.

Step 2: Choose edge 0-1 with weight 1, add node 1 to MST

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

Why: Edge 0-1 is the cheapest edge connecting a new node to MST.

Step 3: Choose edge 1-2 with weight 2, add node 2 to MST

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

Why: Edge 1-2 is now the cheapest edge connecting a new node to MST.

Step 4: Choose edge 2-3 with weight 3, add node 3 to MST

MST nodes: [0,1,2,3]
All nodes connected

Why: Edge 2-3 is the cheapest edge connecting the last node to MST.

Result:

MST edges: 0-1(1) -> 1-2(2) -> 2-3(3)
Total weight: 6

Annotated Code

DSA Go

package main

import (
	"container/heap"
	"fmt"
)

type Edge struct {
	dest, weight int
}

type Graph [][]Edge

// Priority queue item
// Holds node and edge weight to reach it
// Implements heap.Interface

type Item struct {
	node, weight int
}

type PriorityQueue []Item

func (pq PriorityQueue) Len() int { return len(pq) }
func (pq PriorityQueue) Less(i, j int) bool { return pq[i].weight < pq[j].weight }
func (pq PriorityQueue) Swap(i, j int) { pq[i], pq[j] = pq[j], pq[i] }

func (pq *PriorityQueue) Push(x interface{}) {
	*pq = append(*pq, x.(Item))
}

func (pq *PriorityQueue) Pop() interface{} {
	n := len(*pq)
	item := (*pq)[n-1]
	*pq = (*pq)[:n-1]
	return item
}

func primMST(graph Graph, start int) []string {
	n := len(graph)
	inMST := make([]bool, n) // track nodes in MST
	pq := &PriorityQueue{}
	heap.Init(pq)

	heap.Push(pq, Item{node: start, weight: 0})

	var mstEdges []string

	for pq.Len() > 0 {
		item := heap.Pop(pq).(Item)
		if inMST[item.node] {
			continue // skip if already in MST
		}

		inMST[item.node] = true

		// Record edge if not start node
		if item.weight != 0 {
			mstEdges = append(mstEdges, fmt.Sprintf("Edge to node %d with weight %d", item.node, item.weight))
		}

		// Add edges from current node
		for _, e := range graph[item.node] {
			if !inMST[e.dest] {
				heap.Push(pq, Item{node: e.dest, weight: e.weight})
			}
		}
	}

	return mstEdges
}

func main() {
	// Graph adjacency list
	// Node 0: edges to 1(1), 2(4)
	// Node 1: edges to 0(1), 2(2), 3(6)
	// Node 2: edges to 0(4), 1(2), 3(3)
	// Node 3: edges to 1(6), 2(3)
	graph := Graph{
		{{1, 1}, {2, 4}},
		{{0, 1}, {2, 2}, {3, 6}},
		{{0, 4}, {1, 2}, {3, 3}},
		{{1, 6}, {2, 3}},
	}

	mst := primMST(graph, 0)
	for _, edge := range mst {
		fmt.Println(edge)
	}
}

heap.Push(pq, Item{node: start, weight: 0})

Initialize priority queue with start node and zero weight

item := heap.Pop(pq).(Item)

Extract node with smallest edge weight to MST

if inMST[item.node] { continue }

Skip nodes already included in MST to avoid cycles

inMST[item.node] = true

Mark current node as included in MST

if item.weight != 0 { mstEdges = append(mstEdges, ... ) }

Record the edge that added this node to MST, skip start node

for _, e := range graph[item.node] { if !inMST[e.dest] { heap.Push(pq, Item{node: e.dest, weight: e.weight}) } }

Add all edges from current node to priority queue if destination not in MST

Output

Edge to node 1 with weight 1

Edge to node 2 with weight 2

Edge to node 3 with weight 3

Complexity Analysis

Time: O(E log V) because each edge is pushed to the priority queue and heap operations take log V time

Space: O(V + E) to store the graph and priority queue

vs Alternative: Compared to Kruskal's algorithm which sorts edges first, Prim's uses a priority queue and can be more efficient on dense graphs

Edge Cases

Empty graph (no nodes)

Returns empty MST without errors

DSA Go

for pq.Len() > 0 { ... } handles empty queue gracefully

Graph with single node and no edges

MST contains only the start node with no edges

DSA Go

if item.weight != 0 { ... } skips adding edges for start node

Disconnected graph

MST includes only nodes reachable from start node

DSA Go

inMST array prevents adding unreachable nodes

When to Use This Pattern

When you need to connect all points with minimum total cost and can grow the solution from any starting point, use Prim's algorithm with a priority queue to pick the cheapest next edge.

Common Mistakes

Mistake: Not checking if a node is already in MST before adding edges, causing cycles or duplicates
Fix: Use a boolean array to track included nodes and skip them when popping from the priority queue

Mistake: Adding edges to MST output for the start node with zero weight
Fix: Skip recording edges when the weight is zero, which only applies to the start node

Summary

Builds a minimum cost tree connecting all nodes by always adding the cheapest edge to a new node.

Use when you want a minimum spanning tree starting from any node and can access edges from nodes efficiently.

The key insight is to grow the MST step-by-step by choosing the smallest edge connecting the tree to a new node.

Practice

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

easy

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

B. To find the shortest path between two vertices

C. To detect cycles in a graph

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 the smallest sum of edge weights.

2. Step 2: Compare with other graph algorithms

   Shortest path and cycle detection are different problems; Prim's focuses on minimum spanning tree.

3. Final Answer:

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

4. Quick Check:

   Prim's Algorithm = Minimum Spanning Tree [OK]

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

Common Mistakes:

• Confusing Prim's with Dijkstra's algorithm
• Thinking it finds shortest paths
• Assuming it detects cycles

2. Which of the following Go code snippets correctly initializes a priority queue for Prim's Algorithm?

easy

A. pq := new(PriorityQueue); pq = nil

B. pq := PriorityQueue{}; pq.Init()

C. pq := make(PriorityQueue, 0); heap.Init(&pq)

D. pq := []int{}; sort.Ints(pq)

Solution

1. Step 1: Check correct priority queue initialization in Go

   PriorityQueue is usually a slice; make with length 0 and call heap.Init(&pq) to initialize.

2. Step 2: Analyze other options

   pq := PriorityQueue{}; pq.Init() misses heap.Init call; C sets pq to nil; D uses slice but sorts, not heap.

3. Final Answer:

   pq := make(PriorityQueue, 0); heap.Init(&pq) -> Option C

4. Quick Check:

   Use make + heap.Init for priority queue [OK]

Hint: Use make and heap.Init to prepare priority queue [OK]

Common Mistakes:

• Forgetting to call heap.Init
• Setting priority queue to nil
• Sorting slice instead of using heap

3. Given the following Go code snippet implementing Prim's Algorithm on a graph with edges:

edges := [][]int{{0,1,2},{1,2,3},{0,2,1}}
// Graph with 3 vertices and edges with weights
// What is the total weight of the minimum spanning tree printed by this code?
func primMST() int {
  n := 3
  graph := make([][]Edge, n)
  for _, e := range edges {
    graph[e[0]] = append(graph[e[0]], Edge{to: e[1], weight: e[2]})
    graph[e[1]] = append(graph[e[1]], Edge{to: e[0], weight: e[2]})
  }
  visited := make([]bool, n)
  pq := &PriorityQueue{}
  heap.Init(pq)
  heap.Push(pq, &Item{vertex: 0, weight: 0})
  total := 0
  for pq.Len() > 0 {
    item := heap.Pop(pq).(*Item)
    if visited[item.vertex] {
      continue
    }
    visited[item.vertex] = true
    total += item.weight
    for _, e := range graph[item.vertex] {
      if !visited[e.to] {
        heap.Push(pq, &Item{vertex: e.to, weight: e.weight})
      }
    }
  }
  return total
}

medium

A. 3

B. 4

C. 5

D. 6

Solution

1. Step 1: Build the graph and list edges

   Edges: (0-1)=2, (1-2)=3, (0-2)=1. Start from vertex 0 with weight 0.

2. Step 2: Run Prim's Algorithm step-by-step

   Pick 0 (weight 0), add edges (0-1:2), (0-2:1). Next pick edge with smallest weight: (0-2:1), total=1. Then pick (0-1:2), total=3. Vertex 1 and 2 visited, done.

3. Final Answer:

   3 -> Option A

4. Quick Check:

   Sum edges 1 + 2 = 3 [OK]

Hint: Pick smallest edges connecting new vertices first [OK]

Common Mistakes:

• Adding all edges instead of minimum spanning tree
• Starting from wrong vertex
• Counting edges twice

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

func primBuggy() int {
  n := 3
  graph := make([][]Edge, n)
  // edges omitted for brevity
  visited := make([]bool, n)
  pq := &PriorityQueue{}
  heap.Init(pq)
  heap.Push(pq, &Item{vertex: 0, weight: 0})
  total := 0
  for pq.Len() > 0 {
    item := heap.Pop(pq).(*Item)
    if visited[item.vertex] {
      total += item.weight
      continue
    }
    visited[item.vertex] = true
    for _, e := range graph[item.vertex] {
      if !visited[e.to] {
        heap.Push(pq, &Item{vertex: e.to, weight: e.weight})
      }
    }
  }
  return total
}

medium

A. Not marking vertices as visited

B. Adding edge weight after skipping visited vertex causes incorrect total

C. Not initializing the priority queue with heap.Init

D. Pushing edges with wrong weight values

Solution

1. Step 1: Identify where total is updated

   Total is increased only when visited[item.vertex] is true, which means vertex already visited.

2. Step 2: Understand why this is wrong

   Weights should be added only when vertex is visited first time, not when skipping duplicates.

3. Final Answer:

   Adding edge weight after skipping visited vertex causes incorrect total -> Option B

4. Quick Check:

   Update total only on first visit [OK]

Hint: Add weight only when vertex visited first time [OK]

Common Mistakes:

• Adding weights multiple times for same vertex
• Forgetting to call heap.Init
• Not marking visited vertices

5. You have a weighted undirected graph with 5 vertices and edges:

(0-1, 2), (0-3, 6), (1-2, 3), (1-3, 8), (1-4, 5), (2-4, 7), (3-4, 9)

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

hard

A. 21

B. 19

C. 20

D. 16

Solution

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

   Start at 0: edges (0-1:2), (0-3:6). Pick (0-1:2), total=2.

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

   Visited: 0,1. Edges: (1-2:3), (1-3:8), (1-4:5), (0-3:6). Pick (1-2:3), total=5.

3. Step 3: Continue adding minimum edges connecting new vertices

   Visited: 0,1,2. Edges: (1-3:8), (1-4:5), (0-3:6), (2-4:7). Pick (1-4:5), total=10.

4. Step 4: Add edges from new vertex 4

   Visited: 0,1,2,4. Edges: (1-3:8), (0-3:6), (2-4:7), (3-4:9). Pick (0-3:6), total=16.

5. Step 5: Add last vertex 3

   Visited all vertices. Total weight = 16.

6. Final Answer:

   16 -> Option D

7. Quick Check:

   Sum edges 2+3+5+6 = 16 [OK]

Hint: Always pick smallest edge connecting visited to unvisited [OK]

Common Mistakes:

• Choosing edges with larger weights first
• Not updating visited vertices correctly
• Counting edges twice