dijkstra.py

dijkstra.py

#!/usr/bin/env python3
import sys

graph = {
    'A': {'B': 4, 'D': 1, 'C': 7},
    'B': {'C': 2},
    'C': {'E': 2, 'F': 3},
    'D': {'B': 2, 'E': 4},
    'E': {'D': 3, 'B': 7},
    'F': {'E': 1},
}

def dijkstra(graph, src, dst):
    shortest_distance = {}      # to store the smallest cost from 'src' to each node.
    predecessor = {}            # to store the link parent <-> child node.
    unseen_nodes = graph        # just another variable to store the list of node we haven't checked in.
    infinity = sys.maxsize      # the infinity (very big number) to tell no DIRECT path from x to y
    path = []                   # to store the shorted path from 'src' to 'dst'

    # Initialize, set all node's cost to infinity
    for node in unseen_nodes:
        shortest_distance[node] = infinity
    
    # however, the first node's cost is always 0
    shortest_distance[src] = 0

    # the actual dijkstra algorithm now starts
    while unseen_nodes:
        min_node = None
        for node in unseen_nodes:
            if min_node is None:
                min_node = node
            elif shortest_distance[node] < shortest_distance[min_node]:
                min_node = node

        for child_node, cost in graph[min_node].items():
            if shortest_distance[min_node] + cost < shortest_distance[child_node]:
                shortest_distance[child_node] = shortest_distance[min_node] + cost
                predecessor[child_node] = min_node  
  
        # remove the node from unseen_nodes list as it is already been checked-in
        unseen_nodes.pop(min_node)

    # print the shortest distance from 'src' to each node
    # print(shortest_distance)

    # bonus: get the shortest path
    current_node = dst
    while current_node != src:
        try:
            path.insert(0, current_node)
            current_node = predecessor[current_node]
        except KeyError:
            print('Path not reachable')
            break
    
    # now print the result
    if shortest_distance[dst] != infinity:
        print("Shortest distance from {} to {} is: {}".format(src, dst, shortest_distance[dst]))
        path.insert(0, src)
        print("Path: {}".format(path))
    else:
        print("There is no path from {} to {}".format(src, dst))

dijkstra(graph, 'A', 'F')