blog/content/posts/2024-08-17-kd-tree-revisited/index.md

113 lines
3.7 KiB
Markdown

---
title: "Kd Tree Revisited"
date: 2024-08-17T14:20:22+01:00
draft: false # I don't care for draft mode, git has branches for that
description: "Simplifying the nearest neighbour search"
tags:
- algorithms
- data structures
- python
categories:
- programming
series:
- Cool algorithms
favorite: false
disable_feed: false
---
After giving it a bit of thought, I've found a way to simplify the nearest
neighbour search (i.e: the `closest` method) for the `KdTree` I implemented in
[my previous post]({{< relref "../2024-08-10-kd-tree/index.md" >}}).
<!--more-->
## The improvement
That post implemented the nearest neighbour search by keeping track of the
tree's boundaries (through `AABB`), and each of its sub-trees (through
`AABB.split`), and testing for the early exit condition by computing the
distance of the search's origin to each sub-tree's boundaries.
Instead of _explicitly_ keeping track of each sub-tree's boundaries, we can
implicitly compute it when recursing down the tree.
To check for the distance between the queried point and the splitting plane of
inner nodes: we simply need to project the origin onto that plane, thus giving
us a minimal bound on the distance of the points stored on the other side.
This can be easily computed from the `axis` and `mid` values which are stored in
the inner nodes: to project the node on the plane we simply replace its
coordinate for this axis by `mid`.
## Simplified search
With that out of the way, let's now see how `closest` can be implemented without
needing to track the tree's `AABB` at the root:
```python
# Wrapper type for closest points, ordered by `distance`
@dataclasses.dataclass(order=True)
class ClosestPoint[T](NamedTuple):
point: Point = field(compare=False)
value: T = field(compare=False)
distance: float
class KdTree[T]:
def closest(self, point: Point, n: int = 1) -> list[ClosestPoint[T]]:
assert n > 0
res = MaxHeap()
# Instead of passing an `AABB`, we give an initial projection point,
# the query origin itself (since we haven't visited any split node yet)
self._root.closest(point, res, n, point)
return sorted(res)
class KdNode[T]:
def closest(
self,
point: Point,
out: MaxHeap[ClosestPoint[T]],
n: int,
projection: Point,
) -> None:
# Same implementation
self.inner.closest(point, out, n, bounds)
class KdLeafNode[T]:
def closest(
self,
point: Point,
out: MaxHeap[ClosestPoint[T]],
n: int,
projection: Point,
) -> None:
# Same implementation
for p, val in self.points.items():
item = ClosestPoint(p, val, dist(p, point))
if len(out) < n:
out.push(item)
elif out.peek().distance > item.distance:
out.pushpop(item)
class KdSplitNode[T]:
def closest(
self,
point: Point,
out: list[ClosestPoint[T]],
n: int,
projection: Point,
) -> None:
index = self._index(point)
self.children[index].closest(point, out, n, projection)
# Project onto the splitting plane, for a minimum distance to its points
projection = projection.replace(self.axis, self.mid)
# If we're at capacity and can't possibly find any closer points, exit
if len(out) == n and dist(point, projection) > out.peek().distance:
return
# Otherwise recurse on the other side to check for nearer neighbours
self.children[1 - index].closest(point, out, n, projection)
```
As you can see, the main difference is in `KdSplitNode`'s implementation, where
we can quickly compute the minimum distance between the search's origin and all
potential points in that subspace.