Dijkstra Shortest Path Algorithm

Python / Other

# https://en.wikipedia.org/wiki/Dijkstra%27s_algorithm
from collections import deque, namedtuple

# 1. mark all nodes unvisited
# 2. set initial node distance to zero, other nodes to infinite
# 3. select unvisited node with smallest distance
# 4. find unvisited neighbors (to selected node), compare distances, and save smallest one
# 5. mark current node visited
# 6. stop if destination is visited, otherwise repeat step 3-6

inf = float('inf')
edge = namedtuple('Edge', 'start, end, cost')

# make_edge
def make_edge(start, end, cost=1):
    return edge(start, end, cost)

class Graph:
    def __init__(self, edges):
        check_edges = [i for i in edges if len(i) not in [2, 3]]
        if check_edges:
            raise ValueError('Error: check_edges, {}'.format(check_edges))

        self.edges = [make_edge(*edge) for edge in edges]

    # vertices
    @property
    def vertices(self):
        return set(
            # turn ([1,2], [3,4]) into [1, 2, 3, 4]
            sum(([edge.start, edge.end] for edge in self.edges), [])
        )

    # get_node_pairs
    def get_node_pairs(self, node_1, node_2, both_ends=True):
        if both_ends:
            node_pairs = [[node_1, node_2], [node_2, node_1]]
        else:
            node_pairs = [[node_1, node_2]]

        return node_pairs

    # remove_edge
    def remove_edge(self, node_1, node_2, both_ends=True):
        node_pairs = self.get_node_pairs(node_1, node_2, both_ends)
        edges = self.edges[:]

        for edge in edges:
            if [edge.start, edge.end] in node_pairs:
                self.edges.remove(edge)

    # add_edge
    def add_edge(self, node_1, node_2, cost=1, both_ends=True):
        node_pairs = self.get_node_pairs(node_1, node_2, both_ends)
        
        for edge in self.edges:
            if [edge.start, edge.end] in node_pairs:
                return ValueError('Error: add_edge, {} {}'.format(node_1, node_2))

        self.edges.append(edge(start=node_1, end=node_2, cost=cost))
        if both_ends:
            self.edges.append(edge(start=node_2, end=node_1, cost=cost))

    # neighbors
    @property
    def neighbors(self):
        neighbors = { vertex: set() for vertex in self.vertices }

        for edge in self.edges:
            neighbors[edge.start].add((edge.end, edge.cost))

        return neighbors

    # dijkstra
    def dijkstra(self, source, destination):
        assert source in self.vertices, 'assert: source do not exist'

        # 1. mark all nodes unvisited
        # 2. set initial node distance to zero, other nodes to infinite
        distances = { vertex: inf for vertex in self.vertices }
        previous_vertices = { vertex: None for vertex in self.vertices }
        distances[source] = 0
        vertices = self.vertices.copy()

        while vertices:
            # 3. select unvisited node with smallest distance
            current_vertex = min(vertices, key = lambda vertex: distances[vertex])

            # 6. stop if destination is visited, otherwise repeat step 3-6
            if distances[current_vertex] == inf:
                break

            # 4. find unvisited neighbors (to selected node), compare distances, and save smallest one
            for neighbor, cost in self.neighbors[current_vertex]:
                alternative_route = distances[current_vertex] + cost

                if (alternative_route < distances[neighbor]):
                    distances[neighbor] = alternative_route
                    previous_vertices[neighbor] = current_vertex

            # 5. mark current node visited
            vertices.remove(current_vertex)

        path, current_vertex = deque(), destination

        while previous_vertices[current_vertex] is not None:
            path.appendleft(current_vertex)
            current_vertex = previous_vertices[current_vertex]

        if path:
            path.appendleft(current_vertex)

        return path

# create example graph
graph = Graph([
    ('a', 'b', 7),
    ('a', 'c', 9),
    ('a', 'f', 14),
    ('b', 'c', 10),
    ('b', 'd', 15),
    ('c', 'd', 11),
    ('c', 'f', 2),
    ('d', 'e', 6),
    ('e', 'f', 9)
])

print(graph.dijkstra('a', 'f'))
# deque(['a', 'c', 'f'])