Python Implementation of Prims Minimum Spanning Tree

Python Implementation of Prims Minimum Spanning Tree

Python Implementation of Prim's Minimum Spanning Tree

    A minimum spanning tree of a graph is a sub-graph that connects all vertices in the graph with a
minimum total weight for the edges. Each edge between the vertices has a weight corresponding
to it and your goal is to connect every vertex while minimizing the total edge weight. Graphs can
have more than one minimum spanning tree. Below is an example of a graph with 5 vertices and
weighted edges that we will be running Prim's algorithm on.

Example graph with edge weights :-

    Prim's algorithm :-  Prim's algorithm is a classic greedy algorithm for finding the MST of a
graph. The general algorithm is :-
  • Create an empty tree M, which will act as a MST.
  • Choose a random vertex v, from the graph.
  • Add the edges that are connected to v into some data structure E.
  • Find the edge in E with the minimum weight, and add this edge to M. Now,make v equal to the other vertex in the edge and repeat from step 3.

    This algorithm runs until the number of edges in MST is equal to the number of vertices in the
graph minus 1. So in the example below, the number of vertices in the graph is 6, so Prim's
algorithm will run until the MST contains 5 edges. Once the algorithm is complete, the MST will
have successfully connected all vertices in the graph with the minimum weighted edges.

Example of Prim's algorithm :-

    There are a few different options we have in what data structure we decide to use to represent
the graph and how we decide to store all the edges. In this implementation we'll represent the
graph as an adjacency matrix, and we have two common options in how we can store the edges :-
  • Store the edges in an array and search through it each time to find the edge with the smallest weight.
  • Store the edges in a binary heap which improves the running time because edges can be found faster.

The implementation below uses method (1). It first converts the graph into a matrix and then on
each iteration of Prim's algorithm, it adds edges to an array which is then searched through linearly
to find the edge with the smallest weight. If there is a tie between edge weights, it simply chooses
the first edge it encounters.

Code :-
def createAdjMatrix(V, G):
      adjMatrix = []
      #create N x N matrix filled with 0 edge weights between all vertices
      for i in range(0, V):
          for j in range(0, V):
      #populate adjacency matrix with correct edge weights
      for i in range(0, len(G)):
          adjMatrix[G[i][0]][G[i][1]] = G[i][2]
          adjMatrix[G[i][1]][G[i][0]] = G[i][2]
      return adjMatrix
  def prims(V, G):
      # create adj matrix from graph
      adjMatrix = createAdjMatrix(V, G)
      #arbitrarily choose initial vertex from graph
      vertex = 0
      #initialize empty edges array and empty MST
      MST = []
      edges = []
      visited = []
      minEdge = [None,None,float('inf')]
      #run prims algorithm until we create an MST
      #that contains every vertex from the graph
      while len(MST) != V-1:
          #mark this vertex as visited
          #add each edge to list of potential edges
          for r in range(0, V):
              if adjMatrix[vertex][r] != 0:
          #find edge with the smallest weight to a vertex
          #that has not yet been visited
          for e in range(0, len(edges)):
              if edges[e][2] < minEdge[2] and edges[e][1] not in visited:
                  minEdge = edges[e]
          #remove min weight edge from list of edges
          #push min edge to MST
          #start at new vertex and reset min edge
          vertex = minEdge[1]
          minEdge = [None,None,float('inf')]
      return MST
  #graph vertices are actually represented as numbers
  #like so: 0, 1, 2, ... V-1
  a, b, c, d, e, f = 0, 1, 2, 3, 4, 5
  #graph edges with weights
  #diagram of graph is shown above
  graph = [[a,b,2],
  #pass the # of vertices and the graph to run prims algorithm 
  print(prims(6, graph))

  Output :-
[[0, 1, 2], [0, 2, 3], [1, 3, 3], [3, 4, 2], [3, 5, 3]]
    Running Time  :-

Because we are using implementation (1) to store the edges and we are representing the graph as
an adjacency matrix, the running time is O(|V|2) where V is the number of vertices the graph
contains. This is because in the worst case, when we add a new vertex to the MST and we store
its edges, the edge with the smallest weight might be at the end of the list requiring us to loop
through the entire array of edges.

This makes the running time V * V = V
2 where the first V represents every vertex in the graph
that is being looped through in the while loop and the second V represents every other vertex that
the first vertex may be connected to in the graph via an edge.