Prim's algorithm for finding the minimum spanning tree

Question

So I was coding Prim's algorithm for practice and this is what I wrote. It works for finding the weight of the minimum spanning tree (MST) but I'm wondering if the loop I am doing to add the the edges in the frontier to the minheap is optimal.

import heapq
from collections import defaultdict
g = defaultdict(list)
weight = 0 
connected = set([])
pq = []
#n is the number of nodes m is the number of edges
n, m = [int(x) for x in raw_input().split(" ")]
#create adjacency list from inputs of the form "vertex1 vertex2 weight"
for _ in xrange(m):
 a, b, w = [int(x) for x in raw_input().split(" ")]
 g[a].append((w, b))
 g[b].append((w, a))
start = int(raw_input())
connected.add(start)
for tup in g[start]:
 heapq.heappush(pq, tup)
while pq:
 w, b = heapq.heappop(pq)
 if b not in connected:
 weight += w
 connected.add(b)
 #by how much does this loop affect the run time of Prims?
 for tup in g[b]:
 heapq.heappush(pq, tup)
print weight

Answer 1

Your while pq loop continues to run long after it needs to. Consider breaking at the bottom of the loop if len(connected) == n, assuming your graph is one component.

Answer 2

#n is the number of nodes m is the number of edges
n, m = [int(x) for x in raw_input().split(" ")] 

well, instead of writing this comment it would be much better to give your variables meaningful and descriptive names, In this way you won't have to read the comment and the variables to know what they are about.

The same applies to g, pq, a.........etc.etc

• vicg: yeah, I understand that for n and m. I only used it because the hackerrank problem used that as its language. pq I agree could be named better to indicate that it's a priorityqueue of the edge weights but I believe using the variable g for naming a generic graph in a graph problem is standard.