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
You can’t perform that action at this time.