Dijkstra Implementation

dijkstra.py

"""
Implementation is inspired by this article:
https://dev.to/mxl/dijkstras-algorithm-in-python-algorithms-for-beginners-dkc

Graph is from this video:
https://www.youtube.com/watch?v=gdmfOwyQlcI
"""

import math

NODES = ["A", "B", "C", "D", "E", "F", "G"]

EDGES = {
    ("A", "B"): 4,
    ("A", "C"): 3,
    ("A", "E"): 7,
    ("B", "C"): 6,
    ("B", "D"): 5,
    ("C", "D"): 11,
    ("C", "E"): 8,
    ("D", "E"): 2,
    ("D", "F"): 2,
    ("D", "G"): 10,
    ("E", "G"): 5,
    ("F", "G"): 3,
}

NEIGHBORS = {n: {} for n in NODES}
for (a, b), weight in EDGES.items():
    NEIGHBORS[a][b] = weight
    NEIGHBORS[b][a] = weight


def dijkstra(source, dest):
    distances = {node: math.inf for node in NODES}
    distances[source] = 0
    route = {node: None for node in NODES}
    unvisited_nodes = list(NODES)

    # Visit all the nodes.
    while unvisited_nodes:
        # Go to the node with the smallest distance
        current_node = min(unvisited_nodes, key=lambda node: distances[node])
        print(f"Visiting {current_node}")

        # Find unvisited neighbors, calculate their distance from current node
        for neighbor, cost in NEIGHBORS[current_node].items():
            alt_route = distances[current_node] + cost

            # If the alt route is shorter than stored distance, use it.
            if alt_route < distances[neighbor]:
                print(f"Found shorter route to {neighbor} via {current_node}")
                distances[neighbor] = alt_route
                route[neighbor] = current_node

        # Drop current from unvisited nodes
        unvisited_nodes.remove(current_node)

    # Start from the dest, use route to backtrace
    path, current = [dest], dest
    while route[current]:
        path.insert(0, route[current])
        current = route[current]

    return path


if __name__ == "__main__":
    path = dijkstra("A", "F")
    print(path)