Understanding Prim’s Algorithm for Minimum Spanning Tree in C++

In the last post, we discussed Minimum Spanning Trees (MST), their applications and, Kruskal’s algorithm to find MST. Here, we’ll discuss Prim’s algorithm for finding the MST.

Prim’s Algorithm

Prim’s Algorithm is a popular greedy algorithm used to find the Minimum Spanning Tree (MST) of a connected, weighted graph. An MST is a subgraph that includes all the vertices of the original graph with the minimum possible total edge weight.
In this blog post, we will dive into the details of Prim’s algorithm and provide a C++ implementation to find the MST of a graph.

Dry Run:

Explanation

1. Initialize:
• Create an empty MST set to store the edges of the MST.
• Choose a starting vertex arbitrarily and mark it as visited.
• Initialize a priority queue (min-heap) to store the edges connected to the visited vertices, prioritized by their edge weights.
2. Iterate until all vertices are visited:
• Select the edge with the minimum weight from the priority queue.
• If the selected edge connects a visited vertex to an unvisited vertex, add it to the MST set and mark the unvisited vertex as visited.
• Add all the edges connected to the newly visited vertex to the priority queue.
3. After visiting all vertices, the MST set will contain all the edges of the MST.

Pseudocode

Prim(G, s)
for each v ∈ G[vertices]:
key[v] = INF
key[s] = 0

Enqueue(pq, {s, key[s]})

while !pq.empty():
u = Extract-Min(pq)
MST[u] = true
if v not in MST[] && weight[u, v] < key[v]:
parent[v] = u
key[v] = weight[u, v]
Enqueue(pq, {v, key[v]})

Code Implementation

//
//  main.cpp
//  Prim's Algorithm
//
//  Created by Himanshu on 28/05/23.
//

#include <iostream>
#include <vector>
#include <queue>
#include <limits>

using namespace std;

#define INF numeric_limits<int>::max()
#define N 5

typedef pair<int, int> graphNode; // pair (node, weight)
vector<vector<graphNode>> graph(N);

void printMST(vector<int>& key, vector<int>& parent) {

cout<<"Minimum Spanning Tree Edges:"<<endl<<endl;

cout<<"Edge \tWeight"<<endl;

for (int i = 1; i < N; ++i) {
if (parent[i] != -1) {
cout<<parent[i]<<" - "<<i<<"\t"<<key[i]<<endl;
}
}
}

void primMST(int source, vector<int> &key, vector<int> &parent, vector<bool> &inMST) {

// Custom comparator for the priority queue
auto cmp = [](const graphNode& a, const graphNode& b) {
return a.second > b.second;
};

priority_queue<graphNode, vector<graphNode>, decltype(cmp)> pq(cmp);

pq.push({source, key[source]});

while (!pq.empty()) {
int u = pq.top().first;
pq.pop();

inMST[u] = true;

for (auto edge : graph[u]) {
int v = edge.first;
int weight = edge.second;

if (!inMST[v] && weight < key[v]) {
parent[v] = u;
key[v] = weight;
pq.push({v, key[v]});
}
}
}

}

void initializeGraph() {

// Example graph
// pair {node, weight}
graph[0].push_back({1, 4});
graph[0].push_back({2, 3});

graph[1].push_back({2, 1});
graph[1].push_back({3, 2});

graph[2].push_back({3, 4});

graph[3].push_back({4, 2});

}

int main() {

vector<int> parent(N, -1); // stores the parent of each node in the MST
vector<int> key(N, INF); // stores the minimum weight to reach each node
vector<bool> inMST(N, false); // indicates whether a node is already in the MST
int source = 0; // source node

initializeGraph();

primMST(source, key, parent, inMST);

printMST(key, parent);

return 0;
}

Output

Minimum Spanning Tree Edges:

Edge 	Weight
0 - 1	4
0 - 2	3
1 - 3	2
3 - 4	2

Time complexity: O( V log V + E log V) = O(E log V), where E is the number of edges and V is the number of vertices. If we use a Fibonacci-heap, we could improve the time complexity to O( E + V log V)

Here’s a working code with step-by-step code run: Prim’s Algorithm