simple implementation of dijkstra algorithm

shortest_path.py

"""
Reference: http://alexhwoods.com/dijkstra/
"""

import collections
import math


class Graph:
    def __init__(self):
        self.vertices = set()
        # make a default value for all edges an empty list
        self.edges = collections.defaultdict(list)
        self.weights = {}

    def add_vertex(self, vetex_name: str):
        self.vertices.add(vetex_name)

    def add_edge(self, from_vertex: str, to_vertex: str, distance: int):
        if from_vertex == to_vertex:
            return None # no cycles allowed
        self.edges[from_vertex].append(to_vertex) # add edge to vertex
        self.weights[(from_vertex, to_vertex)] = distance

    def __str__(self):
        return "Vertices: {vertices}\n Edges: {edges}\n Weights: {weights}"\
                .format(
                    vertices=self.vertices,
                    edges=self.edges,
                    weights=self.weights
                )


def dijkstra(graph: Graph, start: str):

    S = set()  # set of viewed vertices
    # dict of shortest paths from start to v;
    # by algorithm length to unvisited v should be infinity
    delta = dict.fromkeys(list(graph.vertices), math.inf)
    # TODO: need to figure out what is needed for
    previous = dict.fromkeys(list(graph.vertices), None)

    # set the path length to start to 0
    delta[start] = 0

    # while there exists a vertex v not in S
    while S != graph.vertices:
        # let v be the closest vertex that has not been visited
        # first iteration will be a `start`
        v = min((set(delta.keys()) - S), key=delta.get)

        # for each neighbor of v not in S
        for neighbor in set(graph.edges[v]) - S:
            # calculate a path
            path_length = graph.weights[v, neighbor] + delta[v]
            if path_length < delta[neighbor]:
                delta[neighbor] = path_length

                # set the previous vertex of neighbor to v
                previous[neighbor] = v

        S.add(v)
    return (delta, previous)


def shortest_path(graph: Graph, start: str, end: str) -> list:
    """
    Uses dijkstra function to output the shortest path
    """
    delta, previous = dijkstra(graph, start)
    path = []
    vertex = end

    while vertex is not None:
        path.append(vertex)
        vertex = previous[vertex]

    path.reverse()
    return path


def main():
    G = Graph()
    G.add_vertex('a')
    G.add_vertex('b')
    G.add_vertex('c')
    G.add_vertex('d')
    G.add_vertex('e')

    G.add_edge('a', 'b', 2)
    G.add_edge('a', 'c', 8)
    G.add_edge('a', 'd', 5)
    G.add_edge('b', 'c', 1)
    G.add_edge('c', 'e', 3)
    G.add_edge('d', 'e', 4)

    # print(G)


    print(dijkstra(G, 'a'))
    print(shortest_path(G, 'a', 'e'))


if __name__ == '__main__':
    main()