Dijkstra algorithm

dijkstra.py

from collections import defaultdict
from math import inf


def find_closest(dist, unvisited):
    # Select node with the least distance to source
    u = None
    min_du = inf
    for v in unvisited:
        if dist[v] < min_du:
            u = v
            min_du = dist[v]
    return u


def dijkstra(graph, source):
    """Shortest path tree and distances from a"""
    unvisited = set(graph.keys())  # Unvisited nodes
    dist = defaultdict(lambda: inf)  # Distances from source to v
    prev = defaultdict(lambda: None)  # Previous node in optimal path from source to v

    dist[source] = 0  # Distance from source to source
    while bool(unvisited):
        u = find_closest(dist, unvisited)
        if u is None:
            break
        unvisited.remove(u)
        for (v, d) in graph[u]:
            if v in unvisited:
                alt = dist[u] + d
                if alt < dist[v]:  # A shorter path to v has been found
                    dist[v] = alt
                    prev[v] = u

    return dist, prev


G = {1: [(2, 7), (3, 9), (6, 14)],
     2: [(3, 10), (4, 15)],
     3: [(4, 11), (6, 2)],
     4: [(5, 6)],
     5: [(6, 9)],
     6: [(5, 9)]}

G2 = {1: [(2, 7)],
      2: [(1, 3)]}

G3 = {1: [(2, 3)],
      2: [(3, 3), (1, 11)],
      3: [(2, 11)]}

print(G)

dists, prevs = dijkstra(G, 1)
print(dict(dists), dict(prevs))