store nodes in a grid

1 files changed, 12 insertions, 3 deletions

diff --git a/store.py b/store.py
index 9e6d4f7..5d81304 100644
--- a/store.py
+++ b/store.py
@@ -1,3 +1,5 @@
+from collections import defaultdict
+from math import floor
 from osm_parser import get_default_parser
 
 
@@ -8,10 +10,15 @@ class Node:
         self.lng = float(lng)
         self.neighbours = []
 
-
     def coord_tuple(self):
         return self.lat, self.lng
 
+    def grid_tuple(self):
+        return floor(self.lat * 1000), floor(self.lng * 1000)
+
+    def __repr__(self):
+        return f"({self.lat}, {self.lng})"
+
 
 parser = None  # Have a global reusable parser object
 
@@ -37,6 +44,7 @@ def extract_osm_nodes(f_name):
     global parser
     parser = get_default_parser(f_name)
     nodes = dict()
+    grid = defaultdict(list)
 
     for node in parser.iter_nodes():
         nodes[node['id']] = Node(node['id'], node['lat'], node['lon'])
@@ -47,12 +55,13 @@ def extract_osm_nodes(f_name):
     for node_id, node in nodes.copy().items():
         if not node.neighbours:
             del nodes[node_id]
+        else:
+            grid[(floor(node.lat * 1000), floor(node.lng * 1000))].append(node_id)
 
-    return nodes
+    return nodes, grid
 
 
 def select_nodes_in_rectangle(nodes, min_lat, max_lat, min_long, max_long):
     return [node for node in nodes.values()
             if min_lat <= node.lat <= max_lat
             and min_long <= node.lng <= max_long]
-