about summary refs log tree commit diff
path: root/scratch/facebook/dijkstras.py
blob: 7031701994a7c50f6eb566ea4310de728236d6cf (plain) (blame)
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
from heapq import heappush, heappop
import random

# Dijkstra's algorithm will traverse a directed graph with weighted edges. If
# the edges aren't weighted, we can pretend that each edges weighs 1. The
# algorithm will find the shortest path between points A and B.

def dijkstra(a, b, graph):
    h = []
    seen = set()
    heappush(h, (0, a, [a], []))
    while h:
        km, x, path, steps = heappop(h)

        if x == b:
            for a, b, d in steps:
                print("{} -> {} => {}".format(a, b, d))
            return path, km

        seen.add(x)
        for c, dist in graph[x]:
            if c not in seen:
                heappush(h, (km + dist, c, path + [c], steps + [(x, c, dist)]))
    return [], float('inf')

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

beg = random.choice(list(graph.keys()))
end = random.choice(list(graph.keys()))
print("Searching for the shortest path from {} -> {}".format(beg, end))
print(dijkstra(beg, end, graph))