From 74d4aa87e67f81e79d3744183fa2d8c1418c0fdb Mon Sep 17 00:00:00 2001 From: Bruno BELANYI Date: Sat, 10 Aug 2024 16:47:59 +0100 Subject: [PATCH] posts: kd-tree: add search --- content/posts/2024-08-10-kd-tree/index.md | 28 +++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/content/posts/2024-08-10-kd-tree/index.md b/content/posts/2024-08-10-kd-tree/index.md index 3941019..407951a 100644 --- a/content/posts/2024-08-10-kd-tree/index.md +++ b/content/posts/2024-08-10-kd-tree/index.md @@ -201,3 +201,31 @@ class KdNode[T]: self.inner = split_leaf(self.inner, split_axis) return res ``` + +### Searching for a point + +Looking for a given point in the tree look very similar to a _BST_'s search, +each leaf node dividing the space into two sub-spaces, only one of which +contains the point. + +```python +class KdTree[T]: + def lookup(self, point: Point) -> T | None: + # Forward to the root node + return self._root.lookup(point) + +class KdNode[T]: + def lookup(self, point: Point) -> T | None: + # Forward to the wrapped node + return self.inner.lookup(point) + +class KdLeafNode[T]: + def lookup(self, point: Point) -> T | None: + # Simply check whether we've stored the point in this leaf + return self.points.get(point) + +class KdSplitNode[T]: + def lookup(self, point: Point) -> T | None: + # Recurse into the child which contains the point + return self.children[self._index(point)].lookup(point) +```