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

2 Answers 2


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

  • \$\begingroup\$ 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. \$\endgroup\$
    – vicg
    Aug 9, 2016 at 16:08

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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