Skip to content

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
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
# add the new vertex to X
# print the edges of the MST
for edge in T:
print edge

This comment has been minimized.

Copy link

@marsp0 marsp0 commented Jan 3, 2017


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

Do you have any idea on how to improve that ?


This comment has been minimized.

Copy link

@emirpoy 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!


This comment has been minimized.

Copy link

@weihesdlegend weihesdlegend 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