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:

IterationSelected EdgeWeightKey ValuesParent Values
1[0, INF, INF, INF, INF][-1, -1, -1, -1, -1]
20 – 14[0, 4, INF, INF, INF][-1, 0, -1, -1, -1]
30 – 23[0, 4, 3, INF, INF][1, 0, 0, -1, -1]
42 – 34[0, 4, 3, 4, INF][-1, 0, 0, 2, -1]
51 – 32[0, 4, 3, 2, INF][-1, 0, 0, 1, -1]
63 – 42[0, 4, 3, 2, 2][-1, 0, 0, 1, 3]

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
  for each v ∈ adjacent(u):
   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);

    key[source] = 0; // Start with the source node
    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

Leave a Reply

Your email address will not be published. Required fields are marked *