Skip to content

Instantly share code, notes, and snippets.

Embed
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
@marsp0

This comment has been minimized.

Copy link

commented Jan 3, 2017

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

This comment has been minimized.

Copy link

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!

@weihesdlegend

This comment has been minimized.

Copy link

commented Sep 12, 2017

edge = sorted(crossing, key=lambda e:graph[e[0]][e[1]])[0]; Is this very efficient? it's O(|V|log(|V|)) right?
Have you tried binary heaps?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.