1 files changed, 62 insertions, 4 deletions
diff --git a/algorithms.py b/algorithms.py
index 6e01287..66ecff5 100644
--- a/algorithms.py
+++ b/algorithms.py
@@ -16,18 +16,76 @@ def length_haversine(p1, p2):
     return 6372797.560856 * c  # return the distance in meters
 
 
-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'. """
+def neighbours_on_edge(p, steps=1):
+    candidates = set()
+    dy = dx = -steps
+    while dx < steps:
+        candidates.add((dx, dy))
+        dx += 1
+    dx = steps
+    while dy <= steps:
+        candidates.add((dx, dy))
+        dy += 1
+    dy = steps
+    while dx >= -steps:
+        candidates.add((dx, dy))
+        dx -= 1
+    dx = -steps
+    while dy >= -steps:
+        candidates.add((dx, dy))
+        dy -= 1
+    dy = -steps
+    candidates = [(p[0] + dx, p[1] + dy) for dx, dy in candidates]
+    candidates = sorted(list(candidates))
+    return candidates
+
+
+def neighbours_containing_nodes(nodes, grid, p):
+    res = []
+    steps = 0
+    while not res:
+        res = [sq for sq in neighbours_on_edge(p, steps)
+               if sq in grid]
+        steps += 1
+    return res
+
+
+def get_closest_in_square(nodes, node_ids, source_node):
     min_node = None
     min_value = None
 
-    for node_id, node in nodes.items():
+    for node_id in node_ids:
+        node = nodes[node_id]
         length = length_haversine(source_node, node)
         if min_node is None or length < min_value:
             min_node = node_id
             min_value = length
 
+    return min_node, min_value
+
+
+def get_closest_node_id(nodes, grid, source_node):
+    """ Find the closest node using a grid search and return its id. """
+    first_squares = neighbours_containing_nodes(nodes, grid, source_node.grid_tuple())
+
+    min_node = None
+    min_value = None
+    for sq in first_squares:
+        for node_id in grid[sq]:
+            node = nodes[node_id]
+            length = length_haversine(source_node, node)
+            if min_node is None or length < min_value:
+                min_node = node_id
+                min_value = length
+
+    for sq in neighbours_on_edge(nodes[min_node].grid_tuple()):
+        for node_id in grid[sq]:
+            node = nodes[node_id]
+            length = length_haversine(source_node, node)
+            if length < min_value:
+                min_node = node_id
+                min_value = length
+
     return min_node