Minimum Spanning Tree - Prim's Algorithm

 

What is Prim's Algorithm?

Like Kruskal's Algorithm, Prim's Algorithm is also a Greedy Algorithm and is used to find the Minimum Spanning Tree of a graph. In Prim's Algorithm, we will start from a source and then keep on selecting the next vertex which has minimum distance. Also, we avoid going into a cycle by marking the node which we have already processed as 'visited'.

What is the difference between Prim's Algorithm and Kruskal's Algorithm?

In Kruskal's Algorithm, our focus was on edges. But in Prim's algorithm, our focus will be on vertices/nodes of the graph.

Prim's Algorithm [Time Complexity: O(ElogV), Spave Complexity: O(V)]

1. Take any vertex/node of the graph as starting vertex/node and mark its distance as 0. Also, mark the distance of all other nodes as infinity.

2. Push all the vertices in a min-heap.

3. Do till min-heap is not empty:

        4. Extract min from minheap and mark that extracted node as visited.

        5. For all the unvisited adjacent nodes/vertices of the current extracted node:

                    6. check if the 'distance' of adjacent node > the weight of edge from neighbout to current extracted node, if yes, update the distance of adjacent node to edge weight.

Code: 

import java.util.*;

// The GraphNode implementing Comparable interface as we are going to use min-heap on these graph nodes.

class GraphNode implements Comparable<GraphNode>{

    String data;

    // used for prim's algorithm

    int distance;

    // adjacent nodes to current node

    ArrayList<GraphNode> neighbours;

    // weight of edges from currentNode

    HashMap<GraphNode,Integer> weightMap;

    

    // constructor

    GraphNode(String x){

        this.data = x;

        // setting the distance to Integer.MAX_VALUE as required by Prim's algorithm

        this.distance = Integer.MAX_VALUE;

        this.neighbours = new ArrayList<>();

        this.weightMap = new HashMap<>();

    }

    

    // We are going to use min-heap on these graphnodes, so we need to tell the minheap how to compare these graph nodes

    // for doing that we need to override the compareTo method of comparable interface

    @Override

    public int compareTo(GraphNode e){

        return this.distance - e.distance;

    }

    

}

public class Main {

    

    // stores list of all the nodes present in the graph

    ArrayList<GraphNode> nodesList;

    

    // constructor

    Main(){

        this.nodesList = new ArrayList<>();

    }

    

    // method to add an undirected weighted edge from node represented by index i and j in the nodesList

    public void addUndirectedWeightedEdge(int i, int j, int wt){

        GraphNode firstNode = this.nodesList.get(i);

        GraphNode secondNode = this.nodesList.get(j);

        

        firstNode.neighbours.add(secondNode);

        firstNode.weightMap.put(secondNode,wt);

        

        secondNode.neighbours.add(firstNode);

        secondNode.weightMap.put(firstNode,wt);

    }

    

    // method to find Minimum Spanning Tree using Prim's Algorithm

    public void prim(){

        // hashset used to mark noded as visited

        HashSet<GraphNode> visited = new HashSet<>();

        // the minHeap

        PriorityQueue<GraphNode> minHeap = new PriorityQueue<>();

        // setting the source to V0

        GraphNode source = this.nodesList.get(0);

        // setting the surce distance to 0, as required by Prim's algorithm

        source.distance = 0;

        

        // adding all vertices in the nodesList to the minHeap

        minHeap.addAll(this.nodesList);

        

        // this keeps track of the total cost of edges in the Minimum Spanning Tree

        int cost = 0;

        

        // while minHeap is not empty 

        while(!minHeap.isEmpty()){

            // extract min 

            GraphNode currentNode = minHeap.poll();

            // mark it as visited

            visited.add(currentNode);

            // update cost

            cost += currentNode.distance;

            // for all the adjacent nodes of the currentNode, check if the distance of neighbour is greater that weight of edge

            // if yes, update distance of edge as weight of edge

            for(GraphNode neighbour : currentNode.neighbours){

                // distanceToNeighbour is edge weight from currentNode to current neighbour

                int distanceToNeighbour = currentNode.weightMap.get(neighbour);

                // if current neighbour is not visited, check distanc and update

                if(!visited.contains(neighbour) && (neighbour.distance > distanceToNeighbour)){

                    neighbour.distance = distanceToNeighbour;

                    

                    // refresh the minheap

                    minHeap.remove(neighbour);

                    minHeap.add(neighbour);

                }

            }

            System.out.println("Node selected: "+ currentNode.data+" || Weight of edge from selected node: " + currentNode.distance);

        }

        

        System.out.println("Total cost of Minimum Spanning Tree = " + cost);

    }

    

    public static void main(String[] args) throws Exception {

        Main graph = new Main();

        

        // adding vertices/nodes

        for(int i=0;i<4;i++){

            graph.nodesList.add(new GraphNode("V"+i));

        }

        

        // adding undirected, weighted edges 

        graph.addUndirectedWeightedEdge(0,1,10);

        graph.addUndirectedWeightedEdge(1,3,15);

        graph.addUndirectedWeightedEdge(3,2,4);

        graph.addUndirectedWeightedEdge(2,0,6);

        graph.addUndirectedWeightedEdge(3,0,5);

        

        graph.prim();

        

    }

}

Output:

Node selected: V0 || Weight of edge from selected node: 0
Node selected: V3 || Weight of edge from selected node: 5
Node selected: V2 || Weight of edge from selected node: 4
Node selected: V1 || Weight of edge from selected node: 10
Total cost of Minimum Spanning Tree = 19