Skip to content

Instantly share code, notes, and snippets.

@potpath
Forked from econchick/gist:4666413
Last active March 1, 2023 23:40
Show Gist options
  • Save potpath/b1cc6383e1116e895ac2ec891f666888 to your computer and use it in GitHub Desktop.
Save potpath/b1cc6383e1116e895ac2ec891f666888 to your computer and use it in GitHub Desktop.
Python implementation of Dijkstra's Algorithm using heapq
import heapq
from collections import defaultdict
class Graph:
def __init__(self, n):
self.nodes = set(range(n))
self.edges = defaultdict(list)
self.distances = {}
def add_edge(self, from_node, to_node, distance):
self.edges[from_node].append(to_node)
self.distances[from_node, to_node] = distance
def dijkstra(graph, initial):
visited = {initial: 0}
h = [(0, initial)]
path = {}
nodes = set(graph.nodes)
while nodes and h:
current_weight, min_node = heapq.heappop(h)
try:
while min_node not in nodes:
current_weight, min_node = heapq.heappop(h)
except IndexError:
break
nodes.remove(min_node)
for v in graph.edges[min_node]:
weight = current_weight + graph.distances[min_node, v]
if v not in visited or weight < visited[v]:
visited[v] = weight
heapq.heappush(h, (weight, v))
path[v] = min_node
return visited, path
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment