about summary refs log tree commit diff
path: root/scratch/deepmind/part_one/dijkstra.py
diff options
context:
space:
mode:
Diffstat (limited to 'scratch/deepmind/part_one/dijkstra.py')
-rw-r--r--scratch/deepmind/part_one/dijkstra.py26
1 files changed, 26 insertions, 0 deletions
diff --git a/scratch/deepmind/part_one/dijkstra.py b/scratch/deepmind/part_one/dijkstra.py
new file mode 100644
index 000000000000..6975dbe4d1d6
--- /dev/null
+++ b/scratch/deepmind/part_one/dijkstra.py
@@ -0,0 +1,26 @@
+# Doing a practice implementation of Dijkstra's algorithm: a priority-first
+# search.
+from heapq import heappush, heappop
+
+
+class Node(object):
+    def __init__(self, value, children):
+        self.value = value
+        self.children = children
+
+
+def shortest_path(a, b):
+    """Return the shortest path from `a` to `b`."""
+    q = []
+    seen = set()
+    heappush((a.value, a, [a]), q)
+
+    while q:
+        d, node, path = heappop(q)
+        if node == b:
+            return path
+        seen.add(node)
+        for child in node.children:
+            if child not in seen:
+                heappush((d + child.value, child, path + [child]), q)
+    raise Exception("Path between nodes A and B does not exist.")