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())
for tup in g[start]:
    heapq.heappush(pq, tup)
while pq:
    w, b = heapq.heappop(pq)
    if b not in connected:
        weight += w
        #by how much does this loop affect the run time of Prims?
        for tup in g[b]:
            heapq.heappush(pq, tup)

print weight
up vote 2 down vote accepted

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.

#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

  • 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. – vicg Aug 9 '16 at 16:08

Your Answer


By clicking "Post Your Answer", you acknowledge that you have read our updated terms of service, privacy policy and cookie policy, and that your continued use of the website is subject to these policies.

Not the answer you're looking for? Browse other questions tagged or ask your own question.