1"""Chapter 18: Geospatial indexing. Tested with Python 3, standard library only.
2
3Running example: the grocery app's delivery riders in Bengaluru.
4"Find riders within 2 km of this store."
5"""
6import math
7import random
8
9EARTH_RADIUS_KM = 6371.0
10BASE32 = "0123456789bcdefghjkmnpqrstuvwxyz" # geohash alphabet (no a, i, l, o)
11
12
13# ---------- distance ----------
14
15def haversine_km(lat1, lng1, lat2, lng2):
16 """Great-circle distance between two points on Earth, in km."""
17 p1, p2 = math.radians(lat1), math.radians(lat2)
18 dp = p2 - p1
19 dl = math.radians(lng2 - lng1)
20 a = math.sin(dp / 2) ** 2 + math.cos(p1) * math.cos(p2) * math.sin(dl / 2) ** 2
21 return 2 * EARTH_RADIUS_KM * math.asin(math.sqrt(a))
22
23
24# ---------- geohash ----------
25
26def geohash_encode(lat, lng, precision):
27 """Alternate halving longitude and latitude; every 5 bits become one character."""
28 lat_lo, lat_hi, lng_lo, lng_hi = -90.0, 90.0, -180.0, 180.0
29 out, bits, value, even = [], 0, 0, True
30 while len(out) < precision:
31 if even: # longitude bit
32 mid = (lng_lo + lng_hi) / 2
33 if lng >= mid:
34 value, lng_lo = value * 2 + 1, mid
35 else:
36 value, lng_hi = value * 2, mid
37 else: # latitude bit
38 mid = (lat_lo + lat_hi) / 2
39 if lat >= mid:
40 value, lat_lo = value * 2 + 1, mid
41 else:
42 value, lat_hi = value * 2, mid
43 even = not even
44 bits += 1
45 if bits == 5:
46 out.append(BASE32[value])
47 bits, value = 0, 0
48 return "".join(out)
49
50
51def geohash_bounds(gh):
52 """Return (lat_lo, lat_hi, lng_lo, lng_hi) of the cell."""
53 lat_lo, lat_hi, lng_lo, lng_hi = -90.0, 90.0, -180.0, 180.0
54 even = True
55 for ch in gh:
56 v = BASE32.index(ch)
57 for shift in range(4, -1, -1):
58 bit = (v >> shift) & 1
59 if even:
60 mid = (lng_lo + lng_hi) / 2
61 lng_lo, lng_hi = (mid, lng_hi) if bit else (lng_lo, mid)
62 else:
63 mid = (lat_lo + lat_hi) / 2
64 lat_lo, lat_hi = (mid, lat_hi) if bit else (lat_lo, mid)
65 even = not even
66 return lat_lo, lat_hi, lng_lo, lng_hi
67
68
69def cell_size_km(gh, at_lat):
70 lat_lo, lat_hi, lng_lo, lng_hi = geohash_bounds(gh)
71 height = (lat_hi - lat_lo) * 111.2
72 width = (lng_hi - lng_lo) * 111.2 * math.cos(math.radians(at_lat))
73 return width, height
74
75
76def geohash_neighbours(gh):
77 """The cell itself plus its 8 neighbours (ignores poles and the 180° line)."""
78 lat_lo, lat_hi, lng_lo, lng_hi = geohash_bounds(gh)
79 clat, clng = (lat_lo + lat_hi) / 2, (lng_lo + lng_hi) / 2
80 dlat, dlng = lat_hi - lat_lo, lng_hi - lng_lo
81 return {geohash_encode(clat + i * dlat, clng + j * dlng, len(gh))
82 for i in (-1, 0, 1) for j in (-1, 0, 1)}
83
84
85class GeohashIndex:
86 """cell -> {rider: (lat, lng)}; like a GEO set in Redis, kept in a dict."""
87
88 def __init__(self, precision):
89 self.precision = precision
90 self.cells = {}
91 self.where = {} # rider -> cell
92 self.cell_moves = 0
93
94 def update(self, rider, lat, lng):
95 cell = geohash_encode(lat, lng, self.precision)
96 old = self.where.get(rider)
97 if old is not None and old != cell:
98 del self.cells[old][rider]
99 self.cell_moves += 1
100 self.cells.setdefault(cell, {})[rider] = (lat, lng)
101 self.where[rider] = cell
102
103 def nearby(self, lat, lng, radius_km):
104 """Candidates from 9 cells, then the exact distance filter."""
105 found = []
106 for cell in geohash_neighbours(geohash_encode(lat, lng, self.precision)):
107 for rider, (rlat, rlng) in self.cells.get(cell, {}).items():
108 d = haversine_km(lat, lng, rlat, rlng)
109 if d <= radius_km:
110 found.append((round(d, 3), rider))
111 return sorted(found)
112
113 def candidates(self, lat, lng):
114 return sum(len(self.cells.get(c, {}))
115 for c in geohash_neighbours(geohash_encode(lat, lng, self.precision)))
116
117
118# ---------- quadtree ----------
119
120class QuadTree:
121 """Split a square into four when it holds more than `capacity` points."""
122
123 def __init__(self, lat_lo, lat_hi, lng_lo, lng_hi, capacity=8, depth=0):
124 self.box = (lat_lo, lat_hi, lng_lo, lng_hi)
125 self.capacity, self.depth = capacity, depth
126 self.points, self.children = [], None
127
128 def insert(self, lat, lng, item):
129 lat_lo, lat_hi, lng_lo, lng_hi = self.box
130 if not (lat_lo <= lat < lat_hi and lng_lo <= lng < lng_hi):
131 return False
132 if self.children is None:
133 self.points.append((lat, lng, item))
134 if len(self.points) > self.capacity and self.depth < 20:
135 self._split()
136 return True
137 return any(c.insert(lat, lng, item) for c in self.children)
138
139 def _split(self):
140 lat_lo, lat_hi, lng_lo, lng_hi = self.box
141 mlat, mlng = (lat_lo + lat_hi) / 2, (lng_lo + lng_hi) / 2
142 d = self.depth + 1
143 self.children = [QuadTree(lat_lo, mlat, lng_lo, mlng, self.capacity, d),
144 QuadTree(lat_lo, mlat, mlng, lng_hi, self.capacity, d),
145 QuadTree(mlat, lat_hi, lng_lo, mlng, self.capacity, d),
146 QuadTree(mlat, lat_hi, mlng, lng_hi, self.capacity, d)]
147 pts, self.points = self.points, []
148 for p in pts:
149 any(c.insert(*p) for c in self.children)
150
151 def query_box(self, lat_lo, lat_hi, lng_lo, lng_hi, out):
152 a_lo, a_hi, g_lo, g_hi = self.box
153 if lat_hi < a_lo or lat_lo >= a_hi or lng_hi < g_lo or lng_lo >= g_hi:
154 return out # no overlap: skip the whole square
155 for lat, lng, item in self.points:
156 if lat_lo <= lat <= lat_hi and lng_lo <= lng <= lng_hi:
157 out.append((lat, lng, item))
158 for c in self.children or []:
159 c.query_box(lat_lo, lat_hi, lng_lo, lng_hi, out)
160 return out
161
162 def nearby(self, lat, lng, radius_km):
163 dlat = radius_km / 111.2
164 dlng = radius_km / (111.2 * math.cos(math.radians(lat)))
165 box = self.query_box(lat - dlat, lat + dlat, lng - dlng, lng + dlng, [])
166 found = [(round(haversine_km(lat, lng, a, b), 3), item) for a, b, item in box]
167 return sorted(f for f in found if f[0] <= radius_km)
168
169 def leaves(self):
170 if self.children is None:
171 return [self]
172 return [leaf for c in self.children for leaf in c.leaves()]
173
174
175# ---------- Level 4: the bug ----------
176
177def buggy_nearby(index, lat, lng, radius_km):
178 """Only looks in the user's own cell."""
179 cell = geohash_encode(lat, lng, index.precision)
180 return sorted((round(haversine_km(lat, lng, a, b), 3), r)
181 for r, (a, b) in index.cells.get(cell, {}).items()
182 if haversine_km(lat, lng, a, b) <= radius_km)
183
184
185# ---------- checks ----------
186
187def brute_force(riders, lat, lng, radius_km):
188 return sorted((round(haversine_km(lat, lng, a, b), 3), r)
189 for r, (a, b) in riders.items()
190 if haversine_km(lat, lng, a, b) <= radius_km)
191
192
193def run_checks():
194 # distance
195 assert abs(haversine_km(0, 0, 1, 0) - 111.19) < 0.01 # one degree of latitude
196 assert abs(haversine_km(12.97, 0, 12.97, 1) - 108.35) < 0.05 # one degree of longitude at 12.97°N
197
198 # geohash: well-known test vector
199 assert geohash_encode(57.64911, 10.40744, 11) == "u4pruydqqvj"
200 store = (12.9716, 77.5946)
201 gh = geohash_encode(*store, 7)
202 lat_lo, lat_hi, lng_lo, lng_hi = geohash_bounds(gh)
203 assert lat_lo <= store[0] < lat_hi and lng_lo <= store[1] < lng_hi
204 assert geohash_encode(*store, 5) == gh[:5] # shared prefix = same bigger cell
205
206 # cell sizes at this latitude (precision 5, 6, 7)
207 sizes = {p: cell_size_km(geohash_encode(*store, p), store[0]) for p in (5, 6, 7)}
208 w5, h5 = sizes[5]
209 assert 4.7 < w5 < 4.9 and 4.8 < h5 < 4.9
210 w6, h6 = sizes[6]
211 assert 1.1 < w6 < 1.3 and 0.5 < h6 < 0.7
212 assert len(geohash_neighbours(gh)) == 9
213
214 # 5,000 riders spread over about 30 km x 30 km around the store
215 rnd = random.Random(18)
216 riders = {f"r{i}": (store[0] + rnd.uniform(-0.135, 0.135),
217 store[1] + rnd.uniform(-0.14, 0.14)) for i in range(5000)}
218 index = GeohashIndex(precision=5)
219 tree = QuadTree(12.0, 14.0, 77.0, 78.5, capacity=16)
220 for r, (a, b) in riders.items():
221 index.update(r, a, b)
222 tree.insert(a, b, r)
223
224 truth = brute_force(riders, *store, 2.0)
225 assert index.nearby(*store, 2.0) == truth # 9-cell search is exact
226 assert tree.nearby(*store, 2.0) == truth # quadtree too
227 candidates = index.candidates(*store)
228 assert len(truth) < candidates < 5000 # far fewer than scanning everyone
229
230 # edge problem: a rider 50 m away but across a cell border
231 edge_lat = lat_lo # bottom edge of the store's precision-7 cell
232 user = (edge_lat + 0.0002, store[1]) # just inside the cell
233 rider = (edge_lat - 0.0002, store[1]) # just outside, in the cell below
234 small = GeohashIndex(precision=7)
235 small.update("asha", *rider)
236 gap_m = haversine_km(*user, *rider) * 1000
237 assert 40 < gap_m < 50
238 assert buggy_nearby(small, *user, 1.0) == [] # missed
239 assert [r for _, r in small.nearby(*user, 1.0)] == ["asha"] # found with neighbours
240
241 # moving riders: only cell changes touch the index structure
242 mover = GeohashIndex(precision=6)
243 lat, lng = store
244 for _ in range(900): # one hour at one update every 4 s
245 lat += 0.0002 # about 22 m per update, i.e. ~20 km/h north (assumed)
246 mover.update("ravi", lat, lng)
247 assert 25 < mover.cell_moves < 40 # 900 updates, only ~1 in 30 changes cell
248
249 # quadtree adapts: dense areas get small squares
250 leaves = tree.leaves()
251 assert all(len(leaf.points) <= 16 for leaf in leaves)
252
253 # Level 2 arithmetic
254 assert 10_000 / 4 == 2_500 # writes per second
255 circle = math.pi * 2 ** 2
256 nine_cells = 9 * w5 * h5
257 assert round(circle, 1) == 12.6 and 200 < nine_cells < 210
258
259 return {"truth": len(truth), "candidates": candidates, "gap_m": round(gap_m),
260 "moves": mover.cell_moves, "leaves": len(leaves),
261 "sizes": {p: (round(w, 2), round(h, 2)) for p, (w, h) in sizes.items()},
262 "nine_cells": round(nine_cells)}
263
264
265if __name__ == "__main__":
266 out = run_checks()
267 print("All checks pass.", out)