Skip to content

Instantly share code, notes, and snippets.

@kachayev
Last active July 28, 2024 13:10
Show Gist options
  • Save kachayev/5990802 to your computer and use it in GitHub Desktop.
Save kachayev/5990802 to your computer and use it in GitHub Desktop.
Dijkstra shortest path algorithm based on python heapq heap implementation
from collections import defaultdict
from heapq import *
def dijkstra(edges, f, t):
g = defaultdict(list)
for l,r,c in edges:
g[l].append((c,r))
q, seen, mins = [(0,f,())], set(), {f: 0}
while q:
(cost,v1,path) = heappop(q)
if v1 not in seen:
seen.add(v1)
path = (v1, path)
if v1 == t: return (cost, path)
for c, v2 in g.get(v1, ()):
if v2 in seen: continue
prev = mins.get(v2, None)
next = cost + c
if prev is None or next < prev:
mins[v2] = next
heappush(q, (next, v2, path))
return float("inf"), None
if __name__ == "__main__":
edges = [
("A", "B", 7),
("A", "D", 5),
("B", "C", 8),
("B", "D", 9),
("B", "E", 7),
("C", "E", 5),
("D", "E", 15),
("D", "F", 6),
("E", "F", 8),
("E", "G", 9),
("F", "G", 11)
]
print "=== Dijkstra ==="
print edges
print "A -> E:"
print dijkstra(edges, "A", "E")
print "F -> G:"
print dijkstra(edges, "F", "G")
@tonyflow
Copy link

More concise with path reconstruction. The node IDs are represented as integers while the edge weights as floats

from typing import *
from heapq import *


class Dijkstra:
    def __init__(self,
                 graph: Dict[int, Dict[int, float]],
                 origin: int):
        self.graph = graph
        self.edge_to: Dict[int, int] = {}
        self.distances: Dict[int, float] = {vertex: float('inf') for vertex in self.graph}
        self.origin = origin
        self._find(self.origin)

    def _find(self, node: int):

        self.distances[node] = 0

        # Priority queue which stores tuples from distance to node id
        # Distance is the first in the tuple order since it needs to have
        # priority when entries are inserted into the priority queue
        priority_queue: List[(float, int)] = [(node, 0)]

        while priority_queue:
            current_node, current_distance = heappop(priority_queue)

            # If the distance currently recorded at the distances dict is
            # bigger than the one pushed to the pq then we do not need to
            # process this entry
            if current_distance > self.distances[current_node]:
                continue

            for n, weight in self.graph[current_node].items():
                updated_distance = current_distance + weight

                if updated_distance < self.distances[n]:
                    self.distances[n] = updated_distance
                    self.edge_to[n] = current_node
                    heappush(priority_queue, (n, updated_distance))

    def reconstruct_path(self, destination: int) -> List[int]:
        node: int = destination
        path: List[int] = []

        while node != self.origin:
            path.append(node)
            node = self.edge_to[node]

        path.reverse()
        return path


if __name__ == '__main__':
    graph = {
        1: {2: 1, 3: 4},
        2: {1: 1, 3: 2, 4: 5},
        3: {1: 4, 2: 2, 4: 1},
        4: {2: 5, 3: 1}
    }

    dijkstra: Dijkstra = Dijkstra(graph, 1)

    print(dijkstra.distances)
    print(dijkstra.reconstruct_path(4))
    print(dijkstra.reconstruct_path(2))
    print(dijkstra.reconstruct_path(3))

@VeNoMouS
Copy link

VeNoMouS commented Apr 1, 2024

friends don't let friends import * 😄

to quote pep8

Wildcard imports (from module import *) should be avoided, as they make it unclear which names are present in the namespace, confusing both readers and many automated tools. There is one defensible use case for a wildcard import, which is to republish an internal interface as part of a public API (for example, overwriting a pure Python implementation of an interface with the definitions from an optional accelerator module and exactly which definitions will be overwritten isn’t known in advance).

When republishing names this way, the guidelines below regarding public and internal interfaces still apply.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment