Dijkstra

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
from collections import defaultdict
from dataclasses import dataclass, field
import heapq

Cost = int
Node = str

@dataclass
class Edge:
    source: Node
    target: Node
    cost: Cost

@dataclass
class Graph:
    # represent graph as an edge list 
    edges: defaultdict[Node, list[Edge]] = field(
        default_factory=lambda: defaultdict(list)
    )

    def add_edge(self, edge: Edge):
        self.edges[edge.source].append(edge)

    def get_shortest_path(self, start: Node, finish: Node):
        # dijkstra

        costs: dict[Node, Cost] = {
            node: float("infinity") for node in self.edges
        }
        costs[start] = 0
        parents: dict[Node, Node] = {}

        priority_queue = [(0, start)]

        while priority_queue:
            current_distance, current_node = heapq.heappop(priority_queue)
            if current_distance > costs[current_node]:
                continue
            for edge in self.edges[current_node]:
                cost_to_target = current_distance + edge.cost
                if cost_to_target < costs[edge.target]:
                    costs[edge.target] = cost_to_target
                    parents[edge.target] = current_node
                    heapq.heappush(priority_queue, (cost_to_target, edge.target))

        print("parents:", parents)
        print("minimum distances:", costs)

        path = [finish]
        parent = parents.get(finish)
        while parent:
            path.append(parent)
            if parent == start:
                break
            parent = parents.get(parent)

        shortest_path = list(reversed(path))
        print("shortest_path:", shortest_path)
        return shortest_path


def test_dijkstra():
    cities = Graph()
    edges = [
        ("ankara", "istanbul", 6),
        ("ankara", "eskisehir", 2),
        ("eskisehir", "istanbul", 3),
        ("eskisehir", "izmir", 12),
        ("istanbul", "izmir", 8),
    ]

    for start, finish, distance in edges:
        cities.add_edge(Edge(start, finish, distance))
        cities.add_edge(Edge(finish, start, distance))

    shortest_path = cities.get_shortest_path("ankara", "izmir")
    assert shortest_path == ["ankara", "eskisehir", "istanbul", "izmir"]


if __name__ == "__main__":
    test_dijkstra()