perfect-postcode/finder/spatial.py

37 lines
1.3 KiB
Python

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