dijkstra.py

import heapq

INF = float('inf')

def dijkstra(graph, weights, s):
  d = {v: INF for v in graph}
  d[s] = 0
  pi = {s: None}
  queue = []
  heapq.heappush(queue, (0, s))
  while queue:
    w, u = heapq.heappop(queue)
    for v in graph[u]:
      d_temp = w + weights[(u, v)]
      if d[v] > d_temp:
        d[v] = d_temp
        heapq.heappush(queue, (d_temp, v))
        pi[v] = u
  return d, pi


if __name__ == '__main__':
  graph = {'A': ['B', 'C'],
           'B': ['C', 'D'],
           'C': ['B', 'D', 'E'],
           'D': ['E'],
           'E': ['D']}
  weights = {('A', 'B'): 10, ('A', 'C'): 3,
             ('B', 'C'): 1, ('B', 'D'): 2,
             ('C', 'B'): 4, ('C', 'D'): 8, ('C', 'E'): 2,
             ('D', 'E'): 7,
             ('E', 'D'): 9}
  print(dijkstra(graph, weights, 'A'))