k-d Trees & Nearest Neighbour
In short: every balanced tree so far split on one key. A k-d tree splits a set of points on alternating coordinates — x at the root, y at its children, x again below that — always at the median, so each split keeps the two halves even. The payoff is nearest-neighbour search: instead of checking every point, you can prune a whole half-plane the moment its splitting line alone is already farther away than the best point found so far. No distance check is needed inside a pruned half — it's ruled out by one comparison, not by visiting it.
Snapping a map tap to the nearest pin
A maps app has a few thousand saved pins and needs to answer "which pin did the user just tap near?" many times a second while they drag the map. A k-d tree built once over the pins answers each query by pruning, not scanning:
The nearest pin to (9, 2) is (8, 1) — found after visiting only 3 of the
6 pins. The other 3, all on the far side of the root's split line, were
ruled out in one comparison each.
Pruning one half-plane
The k-d tree over 6 pins, alternating x-splits (root, and its grandchildren) and y-splits (the middle level). Searching for the pin nearest (9, 2).
Every node is badged by which coordinate it split on when it was built — x at the root and its grandchildren, y one level down. Binary tree, 6 nodes, height 3.
A sensor grid that only ever splits one way
A warehouse's floor sensors are logged as (x, y) pairs and a robot needs
"closest working sensor to here." The mistake people actually make building
this: hard-coding the split to always use x, forgetting to alternate
with depth.
def axis_for_depth(depth):
return depth % 2 # 0 = x, 1 = y
print(axis_for_depth(0), axis_for_depth(1), axis_for_depth(2), axis_for_depth(3))If every split used axis = 0 instead of depth % 2, points that share
similar x-coordinates but spread widely in y would still end up on one
lopsided branch — the tree degrades toward a 1-D structure built on a
2-D problem, and the height guarantee from alternating axes disappears.
This prints [0] (only ever splits on x) instead of [0, 1]. Fix `axis` to depend on `depth`.
What it costs
| operation | cost | why |
|---|---|---|
| build | O(n log n) | sorting at each of O(log n) levels, O(n) total per level |
| nearest-neighbour search, balanced data | O(log n) expected | each level prunes roughly half the remaining points |
| nearest-neighbour search, worst case | O(n) | a query that can't rule out any branch (e.g. all points on one line) still visits everything |
| space | O(n) | one node per point, no extra structure |
Quick practice
def squared_dist(a, b):
return (a[0]-b[0])**2 + (a[1]-b[1])**2
print(squared_dist((7,2), (9,2)) < squared_dist((8,1), (9,2)))This prints (9, 6), which is squared-distance 17 from (8, 10), instead of the true nearest point (10, 9), squared-distance 5. The near side alone isn't enough — add back the pruning check that decides whether to explore the far side too.
🎯 Your turn
Build nearest_neighbor(points, query): build a k-d tree from points
(alternating x/y splits, always the median), then search it for the point
nearest query, pruning a branch whenever the splitting plane alone is
already farther than the best distance found so far. Return {"nearest": [x, y] or None, "distance": rounded to 4 decimals or None, "visited": ..., "brute_force_visited": len(points)}.
nearest_neighbor([[2,3],[5,4],[9,6],[4,7],[8,1],[7,2]], [9,2])→{"nearest": [8,1], "distance": 1.4142, "visited": 3, "brute_force_visited": 6}
Hint: build once with the median-split recursion from this lesson, then
search once: visit the near side first, and only visit the far side if
diff * diff (the squared distance to the splitting plane) is still less
than the best squared distance found so far.
Press ▶ Run, try the Live App chip, then ✓ Check.