Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
Prim's algorithm, in Python.
from sys import argv
import re
# open the file and get read to read data
file = open(argv[1], "r");
p = re.compile("\d+");
# initialize the graph
vertices, edges = map(int, p.findall(file.readline()))
graph = [[0]*vertices for _ in range(vertices)]
# populate the graph
for i in range(edges):
u, v, weight = map(int, p.findall(file.readline()))
graph[u][v] = weight
graph[v][u] = weight
# initialize the MST and the set X
T = set();
X = set();
# select an arbitrary vertex to begin with
X.add(0);
while len(X) != vertices:
crossing = set();
# for each element x in X, add the edge (x, k) to crossing if
# k is not in X
for x in X:
for k in range(vertices):
if k not in X and graph[x][k] != 0:
crossing.add((x, k))
# find the edge with the smallest weight in crossing
edge = sorted(crossing, key=lambda e:graph[e[0]][e[1]])[0];
# add this edge to T
T.add(edge)
# add the new vertex to X
X.add(edge[1])
# print the edges of the MST
for edge in T:
print edge

Hey,

I am currently reading a book on algorithms and data structures. I am trying to implement Prim's algorithm in Python, but I do not want to use adjacency matrix. The issue is that I need a heap to get logn extraction, but afterwards I need also a structure to get fast access to the edges. something like that :

build a dict from the input (list of vertices and another one for the edges)
build a heap
start iterating over the heap
exctract min weight from the heap
start iterating over its edges to get the next min <---- the problem is here and this is why i use heap AND a dict
set the min weight from the iteration
repeat

Do you have any idea on how to improve that ?

emirpoy commented Feb 20, 2017

It more sounded like a Kruskal Algo than Prim's algo to find MST.
But for your question, i believe priority queue to keep edge weights should work!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment