show stats

1 files changed, 34 insertions, 3 deletions

diff --git a/algorithms.py b/algorithms.py
index ba06641..0537607 100644
--- a/algorithms.py
+++ b/algorithms.py
@@ -16,6 +16,14 @@ def length_haversine(p1, p2):
     return 6372797.560856 * c  # return the distance in meters
 
 
+def path_length(nodes, path):
+    """ Calculate the length of a path of node IDs. """
+    dist = 0
+    for i in range(len(path)-1):
+        dist += length_haversine(nodes[path[i]], nodes[path[i+1]])
+    return dist
+
+
 def get_closest_node_id(nodes, source_node):
     """ Search through all nodes and return the id of the node
     that is closest to 'source_node'. """
@@ -31,8 +39,8 @@ def get_closest_node_id(nodes, source_node):
     return min_node
 
 
-def find_shortest_path(nodes, source_id, target_id):
-    """ Return the shortest path using Dijkstra's algortihm. """
+def find_shortest_path_dijkstra(nodes, source_id, target_id):
+    """ Return the shortest path using Dijkstra's algorithm. """
     # queue contains multiple (walk_dist, (node_0, node_1, ... node_n))-tuples
     # where (node_0, node_1, ... node_n) is a walk to node_n
     # and walk_dist is the total length of the walk in meters
@@ -58,8 +66,31 @@ def find_shortest_path(nodes, source_id, target_id):
                 continue
             # otherwise this MIGHT be the shortest walk to neighbour
             # so put it in the queue
-            new_dist = walk_dist + length_haversine(nodes[walk_end], neighbour) + length_haversine(neighbour, nodes[target_id])
+            new_dist = walk_dist + length_haversine(nodes[walk_end], neighbour)
             new_walk = walk + (neighbour.id,)
             heapq.heappush(queue, (new_dist, new_walk))
     # no path found
     return None
+
+
+def find_shortest_path_astar(nodes, source_id, target_id):
+    """ Return the shortest path using A*. """
+    queue = [(0, (source_id,))]
+    visited = set()
+
+    while queue:
+        walk_dist, walk = heapq.heappop(queue)
+        walk_end = walk[-1]
+        if walk_end == target_id:
+            return walk
+        if walk_end in visited:
+            continue
+        visited.add(walk_end)
+        for neighbour in nodes[walk_end].neighbours:
+            if neighbour in visited:
+                continue
+            # simple heuristic
+            new_dist = walk_dist + length_haversine(nodes[walk_end], neighbour) + length_haversine(neighbour, nodes[target_id])
+            new_walk = walk + (neighbour.id,)
+            heapq.heappush(queue, (new_dist, new_walk))
+    return None