Prim's Algorithm.py

import sys
class Graph :
  def __init__(self,vertexNum,edges,nodes):
    self.edges = edges
    self.nodes = nodes
    self.vertexNum = vertexNum
    self.MST = []
  def printSolution(self):
       print("Edge : Weight ")
       for s,d,w in self.MST :
         print("%s -> %s: %s"%(s,d,w))
  def primsAlgo(self):
       visited = [0]*self.vertexNum
       edgeNum = 0
       visited[0] = True
       while edgeNum < self.vertexNum-1 :
         min = sys.maxsize
         for i in range(self.vertexNum):
           if visited[i]:
             for j in range(self.vertexNum):
               if ((not visited[j]) and self.edges[i][j]):
                 if min > self.edges[i][j]:
                   min = self.edges[i][j]
                   s = i
                   d = j
         self.MST.append([self.nodes[s],self.nodes[d],self.edges[s][d]])
         visited[d] = True
         edgeNum +=1
       self.printSolution()
edges = []
n = int(input("Enter size :"))
print("Enter the entries :")
for i in range(n):
    a = []
    for j in range(n):
        a.append(int(input()))
    edges.append(a)
nodes = list(map(str,input("Enter nodes : ").split()))
g = Graph(len(nodes),edges,nodes)
g.primsAlgo()