Implementing Prim's with a \$O(m\log n)\$ run time using heaps, similar to Dijkstra heap implementation. Where \$n\$ is the number of nodes, and \$m\$ is the number of edges. Using the heapq
library I tried to create a delete function for my heap ('unexplored') so that this heap only stores crossing edges at every iteration of outer While loop. That didn't work, I got amazing runtime but wrong result, so I just cut corners and used heapify
and inf
. But now the run time is around \$O(nm)\$, but I do get the correct answer now though.
graph
is a dict, which contains adjacency lists of an undirected graph as:
{v1: [[v2, cost2], [v3, cost3],... ]... }
from math import inf
from heapq import heappush, heappop, _siftup, _siftdown, heapify
from time import time
def prims(graph):
start = list(graph)[0]
explored, treeCosts = set([start]), []
unexplored = []
for v in graph[start]:
heappush(unexplored, [v[1], v[0]])
while explored != set(graph):
winner = heappop(unexplored)
explored |= set([winner[1]])
treeCosts.append(winner[0])
#deleting edges which have head at 'winner'
for v in unexplored:
if v[1] == winner[1]:
v[0] = inf
'''i = unexplored.index(v)
unexplored[i], unexplored[-1] = unexplored[-1], unexplored[i]
unexplored.pop()
if i < len(unexplored):
_siftup(unexplored, i)
_siftdown(unexplored, 0, i)'''
heapify(unexplored)
#adding the new crossing edges to heap
for v in graph[winner[1]]:
if v[0] not in explored:
heappush(unexplored, [v[1], v[0]])
return sum(treeCosts)
if __name__ == '__main__':
draft, graph = open('undirectedGraph (weighted, 500).txt').read().splitlines(), {}
for line in draft[1:]:
edge = list(map(int, line.split()))
if edge[0] not in graph:
graph[edge[0]] = []
if edge[1] not in graph:
graph[edge[1]] = []
graph[edge[0]].append(edge[1:])
graph[edge[1]].append([edge[0], edge[2]])
startTime = time()
print('Good prims: ' + str(prims(graph)) + ', ' + 'Time: ' + str(time() - startTime))