More

finder/spatial.py

@@ -0,0 +1,33 @@
+import logging
+import math
+from collections import defaultdict
+
+from constants import GRID_CELL_SIZE
+
+log = logging.getLogger("rightmove")
+
+
+class PostcodeSpatialIndex:
+    """Grid-based spatial index over arcgis postcodes for nearest-lookup."""
+
+    def __init__(self, lats: list[float], lngs: list[float], postcodes: list[str]):
+        self.grid: dict[tuple[int, int], list[tuple[float, float, str]]] = defaultdict(list)
+        for lat, lng, pcd in zip(lats, lngs, postcodes):
+            gx = int(math.floor(lng / GRID_CELL_SIZE))
+            gy = int(math.floor(lat / GRID_CELL_SIZE))
+            self.grid[(gx, gy)].append((lat, lng, pcd))
+        log.info("Postcode spatial index: %d cells, %d postcodes", len(self.grid), len(lats))
+
+    def nearest(self, lat: float, lng: float) -> str | None:
+        gx = int(math.floor(lng / GRID_CELL_SIZE))
+        gy = int(math.floor(lat / GRID_CELL_SIZE))
+        best_dist = float("inf")
+        best_pcd = None
+        for dx in range(-1, 2):
+            for dy in range(-1, 2):
+                for plat, plng, pcd in self.grid.get((gx + dx, gy + dy), []):
+                    d = (plat - lat) ** 2 + (plng - lng) ** 2
+                    if d < best_dist:
+                        best_dist = d
+                        best_pcd = pcd
+        return best_pcd