ProgramGuru

Prim's Algorithm for Minimum Spanning Tree

Problem and examples

Given a connected, undirected, and weighted graph with V vertices and E edges, the goal is to find a Minimum Spanning Tree (MST). A Minimum Spanning Tree is a subset of the edges that connects all the vertices with the minimum possible total edge weight and without forming any cycles.

Prim's Algorithm helps us construct this MST by starting from a node and greedily choosing the next lightest edge that connects a visited node to an unvisited node.

Sometimes, the problem may also ask to return the MST edges themselves (i.e., a list of edge pairs {u, v} included in the MST).

Graph (Edge List) Vertices (V) MST Sum Included Edges
[(0,1,2), (0,3,6), (1,2,3), (1,3,8), (1,4,5), (2,4,7), (3,4,9)] 5 16 {0,1}, {1,2}, {1,4}, {0,3}
[(0,1,10), (0,2,6), (0,3,5), (1,3,15), (2,3,4)] 4 19 {0,3}, {2,3}, {0,1}
[] 0 0 []
[(0,1,1)] 2 1 {0,1}

How it works

Understanding the Problem

We are given a connected, weighted, undirected graph. Our task is to find a Minimum Spanning Tree (MST), which is a subset of edges that connects all vertices in the graph with the minimum possible total edge weight, and without forming any cycles.

Prim’s Algorithm helps us solve this problem using a greedy approach. At every step, it picks the smallest weight edge that connects a node already in the MST to a node outside of it. This ensures we keep the tree connected and avoid cycles.

Step-by-Step Solution with Example

Step 1: Choose a starting point

We start from an arbitrary node. Let’s say we begin from node 0.

Step 2: Add all edges from the start node to a Min Heap

This allows us to always access the smallest available edge. Each heap entry stores a tuple like (weight, from_node, to_node).

Step 3: Pop the smallest edge

We take the edge with the least weight from the heap. If the destination node is not yet in the MST, we add it to the MST and include this edge in our result.

Step 4: Add new edges from the newly included node

From the new node, push all edges to unvisited neighbors into the heap. This helps explore other possible minimum connections.

Step 5: Repeat until all nodes are in the MST

We continue this process of picking the smallest edge and expanding the MST until all vertices are included.

Step 6: Example

Let’s consider the following graph with 5 nodes and these edges:


Edges: 
0 --(1)-- 1
0 --(3)-- 2
1 --(1)-- 2
1 --(6)-- 3
2 --(5)-- 3
3 --(2)-- 4

We start at node 0:

  • Add edges (0-1, weight 1) and (0-2, weight 3) to the heap.
  • Pick edge (0-1). Add node 1 to MST. Total weight = 1
  • Add edges (1-2, weight 1), (1-3, weight 6)
  • Pick edge (1-2). Add node 2. Total weight = 2
  • Add edge (2-3, weight 5)
  • Pick edge (2-3). Add node 3. Total weight = 7
  • Add edge (3-4, weight 2)
  • Pick edge (3-4). Add node 4. Total weight = 9

Final MST edges: (0-1), (1-2), (2-3), (3-4)

Total weight of MST: 9

Edge Cases

  • Single node graph: No edges, MST weight is 0.
  • Disconnected graph: Prim’s Algorithm assumes the graph is connected. If not, we cannot form a valid MST.
  • Multiple edges with same weight: The algorithm still works as it always chooses the minimum available weight. Tie-breaking is handled by heap order.
  • Negative weight edges: Prim’s Algorithm can handle them safely, as it doesn’t depend on edge weight signs—only comparisons.

Finally

Prim’s Algorithm is an elegant and efficient way to find a Minimum Spanning Tree using a greedy approach. The use of a Min Heap ensures that we always pick the optimal next edge. For large graphs, the heap makes the algorithm run efficiently in O(E log V) time.

Understanding each step and visualizing the MST growth helps beginners grasp how the algorithm avoids cycles and ensures minimum total weight.

  1. Initialize a minHeap and push the first node (weight 0, node 0).
  2. Create a visited set to track included vertices.
  3. Maintain a totalWeight and optionally an mstEdges list to store the MST edges.
  4. While the heap is not empty and not all vertices are included:
    1. Extract the edge with the smallest weight.
    2. If the destination node is not visited, add it to the MST.
    3. Add the edge’s weight to totalWeight.
    4. Add all adjacent unvisited edges to the heap.
  5. After the loop, totalWeight holds the MST cost and mstEdges (if used) contains the MST structure.

Implementation and efficiency

C implementation

#include <stdio.h>
#include <stdbool.h>
#include <limits.h>

#define V 5

int minKey(int key[], bool mstSet[]) {
    int min = INT_MAX, min_index = -1;
    for (int v = 0; v < V; v++) {
        if (!mstSet[v] && key[v] < min) {
            min = key[v];
            min_index = v;
        }
    }
    return min_index;
}

void primMST(int graph[V][V]) {
    int parent[V];
    int key[V];
    bool mstSet[V];

    for (int i = 0; i < V; i++) {
        key[i] = INT_MAX;
        mstSet[i] = false;
    }

    key[0] = 0;
    parent[0] = -1;

    for (int count = 0; count < V - 1; count++) {
        int u = minKey(key, mstSet);
        mstSet[u] = true;

        for (int v = 0; v < V; v++) {
            if (graph[u][v] && !mstSet[v] && graph[u][v] < key[v]) {
                parent[v] = u;
                key[v] = graph[u][v];
            }
        }
    }

    int totalWeight = 0;
    printf("Edges in MST:\n");
    for (int i = 1; i < V; i++) {
        printf("%d - %d\n", parent[i], i);
        totalWeight += graph[i][parent[i]];
    }
    printf("Total MST Weight: %d\n", totalWeight);
}

int main() {
    int graph[V][V] = {
        {0, 2, 0, 6, 0},
        {2, 0, 3, 8, 5},
        {0, 3, 0, 0, 7},
        {6, 8, 0, 0, 9},
        {0, 5, 7, 9, 0}
    };

    primMST(graph);
    return 0;
}

C++ implementation

#include <iostream>
#include <vector>
#include <climits>
using namespace std;

const int V = 5;

int minKey(int key[], bool mstSet[]) {
    int min = INT_MAX, min_index = -1;
    for (int v = 0; v < V; v++) {
        if (!mstSet[v] && key[v] < min) {
            min = key[v];
            min_index = v;
        }
    }
    return min_index;
}

void primMST(int graph[V][V]) {
    int parent[V];
    int key[V];
    bool mstSet[V];

    for (int i = 0; i < V; i++) {
        key[i] = INT_MAX;
        mstSet[i] = false;
    }

    key[0] = 0;
    parent[0] = -1;

    for (int count = 0; count < V - 1; count++) {
        int u = minKey(key, mstSet);
        mstSet[u] = true;

        for (int v = 0; v < V; v++) {
            if (graph[u][v] && !mstSet[v] && graph[u][v] < key[v]) {
                parent[v] = u;
                key[v] = graph[u][v];
            }
        }
    }

    int totalWeight = 0;
    cout << "Edges in MST:" << endl;
    for (int i = 1; i < V; i++) {
        cout << parent[i] << " - " << i << endl;
        totalWeight += graph[i][parent[i]];
    }
    cout << "Total MST Weight: " << totalWeight << endl;
}

int main() {
    int graph[V][V] = {
        {0, 2, 0, 6, 0},
        {2, 0, 3, 8, 5},
        {0, 3, 0, 0, 7},
        {6, 8, 0, 0, 9},
        {0, 5, 7, 9, 0}
    };

    primMST(graph);
    return 0;
}

Python implementation

def prim_mst(graph):
    V = len(graph)
    key = [float('inf')] * V
    parent = [None] * V
    key[0] = 0
    mst_set = [False] * V

    for _ in range(V):
        u = min((k for k in range(V) if not mst_set[k]), key=lambda x: key[x])
        mst_set[u] = True

        for v in range(V):
            if graph[u][v] and not mst_set[v] and graph[u][v] < key[v]:
                key[v] = graph[u][v]
                parent[v] = u

    total_weight = 0
    print("Edges in MST:")
    for i in range(1, V):
        print(f"{parent[i]} - {i}")
        total_weight += graph[i][parent[i]]
    print(f"Total MST Weight: {total_weight}")

if __name__ == "__main__":
    graph = [
        [0, 2, 0, 6, 0],
        [2, 0, 3, 8, 5],
        [0, 3, 0, 0, 7],
        [6, 8, 0, 0, 9],
        [0, 5, 7, 9, 0]
    ]
    prim_mst(graph)

Java implementation

public class PrimMST {
    private static final int V = 5;

    private int minKey(int[] key, boolean[] mstSet) {
        int min = Integer.MAX_VALUE, minIndex = -1;
        for (int v = 0; v < V; v++) {
            if (!mstSet[v] && key[v] < min) {
                min = key[v];
                minIndex = v;
            }
        }
        return minIndex;
    }

    public void primMST(int[][] graph) {
        int[] parent = new int[V];
        int[] key = new int[V];
        boolean[] mstSet = new boolean[V];

        for (int i = 0; i < V; i++) {
            key[i] = Integer.MAX_VALUE;
            mstSet[i] = false;
        }

        key[0] = 0;
        parent[0] = -1;

        for (int count = 0; count < V - 1; count++) {
            int u = minKey(key, mstSet);
            mstSet[u] = true;

            for (int v = 0; v < V; v++) {
                if (graph[u][v] != 0 && !mstSet[v] && graph[u][v] < key[v]) {
                    parent[v] = u;
                    key[v] = graph[u][v];
                }
            }
        }

        int totalWeight = 0;
        System.out.println("Edges in MST:");
        for (int i = 1; i < V; i++) {
            System.out.println(parent[i] + " - " + i);
            totalWeight += graph[i][parent[i]];
        }
        System.out.println("Total MST Weight: " + totalWeight);
    }

    public static void main(String[] args) {
        int[][] graph = {
            {0, 2, 0, 6, 0},
            {2, 0, 3, 8, 5},
            {0, 3, 0, 0, 7},
            {6, 8, 0, 0, 9},
            {0, 5, 7, 9, 0}
        };

        PrimMST prim = new PrimMST();
        prim.primMST(graph);
    }
}

JavaScript implementation

function primsMST(V, adj) {
  const visited = new Array(V).fill(false);
  const minHeap = [[0, 0]]; // [weight, vertex]
  const mstEdges = [];
  let totalWeight = 0;

  while (minHeap.length) {
    minHeap.sort((a, b) => a[0] - b[0]);
    const [weight, u] = minHeap.shift();

    if (visited[u]) continue;
    visited[u] = true;
    totalWeight += weight;

    for (const [v, wt] of adj[u]) {
      if (!visited[v]) {
        minHeap.push([wt, v]);
        mstEdges.push([u, v]);
      }
    }
  }

  console.log("MST Weight:", totalWeight);
  return { totalWeight, mstEdges };
}

// Example usage
const V = 5;
const adj = [
  [[1,2], [3,6]],
  [[0,2], [2,3], [3,8], [4,5]],
  [[1,3], [4,7]],
  [[0,6], [1,8], [4,9]],
  [[1,5], [2,7], [3,9]]
];

console.log(primsMST(V, adj));

Go implementation

package main

import (
	"fmt"
	"math"
)

const V = 5

func minKey(key []int, mstSet []bool) int {
	min := math.MaxInt64
	minIndex := -1

	for v := 0; v < V; v++ {
		if !mstSet[v] && key[v] < min {
			min = key[v]
			minIndex = v
		}
	}
	return minIndex
}

func primMST(graph [V][V]int) {
	parent := make([]int, V)
	key := make([]int, V)
	mstSet := make([]bool, V)

	for i := 0; i < V; i++ {
		key[i] = math.MaxInt64
		mstSet[i] = false
	}

	key[0] = 0
	parent[0] = -1

	for count := 0; count < V-1; count++ {
		u := minKey(key, mstSet)
		mstSet[u] = true

		for v := 0; v < V; v++ {
			if graph[u][v] != 0 && !mstSet[v] && graph[u][v] < key[v] {
				parent[v] = u
				key[v] = graph[u][v]
			}
		}
	}

	totalWeight := 0
	fmt.Println("Edges in MST:")
	for i := 1; i < V; i++ {
		fmt.Printf("%d - %d\n", parent[i], i)
		totalWeight += graph[i][parent[i]]
	}
	fmt.Println("Total MST Weight:", totalWeight)
}

func main() {
	graph := [V][V]int{
		{0, 2, 0, 6, 0},
		{2, 0, 3, 8, 5},
		{0, 3, 0, 0, 7},
		{6, 8, 0, 0, 9},
		{0, 5, 7, 9, 0},
	}

	primMST(graph)
}

Rust implementation

fn min_key(key: &[i32], mst_set: &[bool]) -> usize {
    let mut min = i32::MAX;
    let mut min_index = 0;

    for (i, &k) in key.iter().enumerate() {
        if !mst_set[i] && k < min {
            min = k;
            min_index = i;
        }
    }
    min_index
}

fn prim_mst(graph: &Vec<Vec<i32>>) {
    let v = graph.len();
    let mut key = vec![i32::MAX; v];
    let mut parent = vec![None; v];
    let mut mst_set = vec![false; v];

    key[0] = 0;
    parent[0] = None;

    for _ in 0..v - 1 {
        let u = min_key(&key, &mst_set);
        mst_set[u] = true;

        for w in 0..v {
            if graph[u][w] != 0 && !mst_set[w] && graph[u][w] < key[w] {
                key[w] = graph[u][w];
                parent[w] = Some(u);
            }
        }
    }

    let mut total_weight = 0;
    println!("Edges in MST:");
    for i in 1..v {
        if let Some(p) = parent[i] {
            println!("{} - {}", p, i);
            total_weight += graph[i][p];
        }
    }
    println!("Total MST Weight: {}", total_weight);
}

fn main() {
    let graph = vec![
        vec![0, 2, 0, 6, 0],
        vec![2, 0, 3, 8, 5],
        vec![0, 3, 0, 0, 7],
        vec![6, 8, 0, 0, 9],
        vec![0, 5, 7, 9, 0],
    ];

    prim_mst(&graph);
}

Kotlin implementation

fun minKey(key: IntArray, mstSet: BooleanArray): Int {
    var min = Int.MAX_VALUE
    var minIndex = -1

    for (v in key.indices) {
        if (!mstSet[v] && key[v] < min) {
            min = key[v]
            minIndex = v
        }
    }
    return minIndex
}

fun primMST(graph: Array<IntArray>) {
    val V = graph.size
    val parent = IntArray(V)
    val key = IntArray(V) { Int.MAX_VALUE }
    val mstSet = BooleanArray(V) { false }

    key[0] = 0
    parent[0] = -1

    for (count in 0 until V - 1) {
        val u = minKey(key, mstSet)
        mstSet[u] = true

        for (v in 0 until V) {
            if (graph[u][v] != 0 && !mstSet[v] && graph[u][v] < key[v]) {
                parent[v] = u
                key[v] = graph[u][v]
            }
        }
    }

    var totalWeight = 0
    println("Edges in MST:")
    for (i in 1 until V) {
        println("${parent[i]} - $i")
        totalWeight += graph[i][parent[i]]
    }
    println("Total MST Weight: $totalWeight")
}

fun main() {
    val graph = arrayOf(
        intArrayOf(0, 2, 0, 6, 0),
        intArrayOf(2, 0, 3, 8, 5),
        intArrayOf(0, 3, 0, 0, 7),
        intArrayOf(6, 8, 0, 0, 9),
        intArrayOf(0, 5, 7, 9, 0)
    )

    primMST(graph)
}

Swift implementation

import Foundation

let V = 5

func minKey(_ key: [Int], _ mstSet: [Bool]) -> Int {
    var min = Int.max
    var minIndex = -1

    for v in 0 ..< V {
        if !mstSet[v] && key[v] < min {
            min = key[v]
            minIndex = v
        }
    }
    return minIndex
}

func primMST(_ graph: [[Int]]) {
    var parent = [Int](repeating: -1, count: V)
    var key = [Int](repeating: Int.max, count: V)
    var mstSet = [Bool](repeating: false, count: V)

    key[0] = 0

    for _ in 0 ..< V - 1 {
        let u = minKey(key, mstSet)
        mstSet[u] = true

        for v in 0 ..< V {
            if graph[u][v] != 0 && !mstSet[v] && graph[u][v] < key[v] {
                parent[v] = u
                key[v] = graph[u][v]
            }
        }
    }

    var totalWeight = 0
    print("Edges in MST:")
    for i in 1 ..< V {
        print("\(parent[i]) - \(i)")
        totalWeight += graph[i][parent[i]]
    }
    print("Total MST Weight: \(totalWeight)")
}

let graph = [
    [0, 2, 0, 6, 0],
    [2, 0, 3, 8, 5],
    [0, 3, 0, 0, 7],
    [6, 8, 0, 0, 9],
    [0, 5, 7, 9, 0]
]

primMST(graph)

TypeScript implementation

function primsMST(V: number, adj: Array<Array<[number, number]>>) {
  const visited: boolean[] = new Array(V).fill(false);
  const minHeap: Array<[number, number]> = [[0, 0]]; // [weight, vertex]
  const mstEdges: Array<[number, number]> = [];
  let totalWeight = 0;

  while (minHeap.length) {
    minHeap.sort((a, b) => a[0] - b[0]);
    const [weight, u] = minHeap.shift()!;

    if (visited[u]) continue;
    visited[u] = true;
    totalWeight += weight;

    for (const [v, wt] of adj[u]) {
      if (!visited[v]) {
        minHeap.push([wt, v]);
        mstEdges.push([u, v]);
      }
    }
  }

  console.log("MST Weight:", totalWeight);
  return { totalWeight, mstEdges };
}

// Example usage
const V = 5;
const adj: Array<Array<[number, number]>> = [
  [[1, 2], [3, 6]],
  [[0, 2], [2, 3], [3, 8], [4, 5]],
  [[1, 3], [4, 7]],
  [[0, 6], [1, 8], [4, 9]],
  [[1, 5], [2, 7], [3, 9]],
];

console.log(primsMST(V, adj));