KD Tree and its usage in robotics

• KD Tree
  • The idea: cut space in half, again and again
  • Toy example: building a tree from 11 points in 3D
  • Nearest neighbor search: go down, then come back up
  • From 1-NN to k-NN: keep the k best
  • Implementation
  • How fast, and when it stops working
• How it is used in robotics for KNN search
  • ICP: finding correspondences
  • Normal estimation and local geometry
  • Filtering: removing outliers
  • Clustering: segmenting obstacles
  • Motion planning: the nearest node in a growing tree
  • Feature matching and place recognition
  • Keeping up with a growing map
  • Libraries
  • Practical pitfalls

Many algorithms in robotics boil down to one question, asked millions of times per second: which stored points are closest to this point? ICP needs the closest map point for every scan point. Normal estimation needs the 20 closest points around every point. A motion planner needs the closest node in its tree. Answering by checking every point costs $O(N)$ per query. With a million map points and a hundred thousand query points per LiDAR scan, that is $10^{11}$ distance computations for a single scan, far too slow for a sensor that sends a new scan every 100 ms.

A k-d tree (k-dimensional tree, Bentley 1975) organizes the points once so that each query only has to look at a handful of them. This post builds one by hand on a small 3D example, walks through nearest-neighbor and k-nearest-neighbor search, and then looks at where it shows up in a robotics stack. For a gentle 2D introduction, the Computerphile video on k-d trees is a good companion.

KD Tree

The idea: cut space in half, again and again

A k-d tree is a binary tree in which every node is one of the data points, and every node also defines a splitting plane:

• Each node splits along one coordinate axis. The axes are used in turn: $x$ at the root, $y$ at depth 1, $z$ at depth 2, then $x$ again, and so on. For a $k$-dimensional point, the axis at depth $d$ is $d \bmod k$.
• The node’s point is the median of its points along that axis. Points with a smaller coordinate go to the left subtree, the rest go to the right subtree.
• The same rule is applied to each half, until no points are left.

Taking the median keeps the tree balanced: each level halves the number of points, so the depth is about $\log_2 N$. For a million points that is only 20 levels.

Each node also owns a region of space: the axis-aligned box carved out by all the splitting planes above it. The root owns all of space, its children own one half each, the grandchildren a quarter, and so on. This is what makes search fast: if a whole box is provably farther away than the best point found so far, every point inside it can be skipped without computing a single distance.

A k-d tree differs from an octree, which cuts every cell into eight equal cubes at its geometric center. A k-d tree cuts through actual data points, one axis at a time, so it adapts to the data. Dense regions get many small cells and empty space gets few large ones. This is a good fit for LiDAR scans, which are very dense close to the sensor and sparse far away.

Toy example: building a tree from 11 points in 3D

Take these 11 points:

Depth 0, split on $x$. Sorted by $x$: C(1), D(2), F(3), B(4), E(5), A(6), I(7), H(8), G(9), K(10), J(11). The median is A, so A is the root and the splitting plane is $x = 6$. The five points with $x < 6$ (C, D, F, B, E) go left, the five with $x > 6$ (I, H, G, K, J) go right.

Depth 1, split on $y$. On the left, sorted by $y$: C(4), D(8), B(9), F(10), E(11). The median B splits at $y = 9$, with {C, D} below and {F, E} above. On the right: H(1), I(3), G(5), K(6), J(7). The median G splits at $y = 5$, with {H, I} below and {K, J} above.

Depth 2, split on $z$. Each group now has two points. For {C, D}, sorted by $z$: D(3), C(6). With an even count we take the upper median, so C becomes the node with plane $z = 6$, and D is its left child. Likewise E splits at $z = 11$ with child F, H splits at $z = 2$ with child I, and J splits at $z = 8$ with child K.

Depth 3. D, F, I and K are alone in their groups, so they are leaves.

                               A (6,2,5)  split x=6
                  x<6 /                            \ x>6
          B (4,9,9)  split y=9                G (9,5,4)  split y=5
         y<9 /          \ y>9                y<5 /           \ y>5
  C (1,4,6) z=6    E (5,11,11) z=11     H (8,1,2) z=2    J (11,7,8) z=8
   z<6 /            z<11 /               z<2 /            z<8 /
  D (2,8,3)        F (3,10,10)          I (7,3,1)        K (10,6,7)

Reading off the regions: B owns the half-space $x < 6$. C owns $x < 6,\ y < 9$. D owns $x < 6,\ y < 9,\ z < 6$, and so on. Every point lies inside the boxes of all its ancestors, which you can check: D = (2, 8, 3) has $x = 2 < 6$, $y = 8 < 9$ and $z = 3 < 6$.

Building costs $O(N \log N)$. There are $\log N$ levels, and each level does $O(N)$ total work if the median is found with a linear-time selection (std::nth_element in C++) instead of a full sort.

Nearest neighbor search: go down, then come back up

Query: which point is closest to $\vec{q} = (0, 8, 12)$?

The search has two phases:

1. Descend. Starting at the root, compare the query with the node along the node’s split axis, and go to the side the query is on. This is the same as an insertion into a binary search tree. It ends in the cell that contains the query, and the points on the way are good first guesses.
2. Backtrack. On the way back up, at each node, ask whether the other side of the splitting plane could contain something closer than the best point so far. If yes, search it too. If no, skip the whole subtree.

The pruning rule. Let the current best squared distance be $d_{\text{best}}^2$, and let the node split on axis $a$ at value $p_a$. Any point $\vec{x}$ on the far side of the plane has $\lvert x_a - q_a \rvert \ge \lvert p_a - q_a \rvert$, because it lies beyond the plane as seen from the query. A squared distance is at least as large as any one of its terms, so

So if $(p_a - q_a)^2 \ge d_{\text{best}}^2$, nothing on the far side can win, and the subtree is skipped. Picture a sphere of radius $d_{\text{best}}$ around the query: the far side only has to be searched if the sphere crosses the plane. The test needs a single subtraction and a multiplication, no matter how many points the subtree holds.

Here are the squared distances from $\vec{q}$ to every point, which is what brute force would compute:

The answer is F, at distance $\sqrt{17} \approx 4.12$. Here is how the tree finds it:

The search computed 5 of the 11 distances. A few things are worth noticing:

• The descent alone gives the wrong answer. It ends at C, which is not even the best point on its own path (B is). The true neighbor F sits just across B’s splitting plane, only 1 unit away in $y$. Backtracking is what makes the search exact.
• Big prunes happen near the root. The last step discards half of the data set, five points, with one comparison. In a balanced tree of a million points, the same test at the root discards half a million.
• Order matters. Searching the near side first makes $d_{\text{best}}$ small early, and a small $d_{\text{best}}$ prunes more.

From 1-NN to k-NN: keep the k best

To find the $k$ nearest neighbors, keep the $k$ best candidates found so far in a max-heap ordered by distance. The top of the heap is the worst of the current $k$, and that is the distance that pruning compares against. Until the heap holds $k$ points, nothing can be pruned, because any point could still make the list.

With $k = 3$ and the same query:

The result is F, B, E. This time D was visited, because with the heap still holding A, the search sphere was huge. A larger $k$ means a larger worst distance, which means fewer prunes, so k-NN is somewhat more expensive than 1-NN. The final prune was also close: 36 against 35.

Radius search (all points within distance $r$) works the same way, with a fixed pruning distance $r^2$ instead of a shrinking one. It returns a variable number of points, which matters for robotics, as we will see.

Implementation

The whole thing fits in about 40 lines of Python:

import heapq
import numpy as np

class Node:
    def __init__(self, idx, axis, left, right):
        self.idx, self.axis, self.left, self.right = idx, axis, left, right

def build(points, ids, depth=0):
    if len(ids) == 0:
        return None
    axis = depth % points.shape[1]                 # cycle x, y, z, x, ...
    ids = ids[np.argsort(points[ids, axis])]
    m = len(ids) // 2                              # median point becomes the node
    return Node(ids[m], axis,
                build(points, ids[:m], depth + 1),
                build(points, ids[m + 1:], depth + 1))

def knn(points, root, q, k):
    heap = []                                      # max-heap of (-dist^2, idx), size <= k

    def search(node):
        if node is None:
            return
        d2 = np.sum((points[node.idx] - q) ** 2)
        if len(heap) < k:
            heapq.heappush(heap, (-d2, node.idx))
        elif d2 < -heap[0][0]:
            heapq.heapreplace(heap, (-d2, node.idx))

        diff = q[node.axis] - points[node.idx, node.axis]
        near, far = (node.left, node.right) if diff < 0 else (node.right, node.left)
        search(near)
        worst = -heap[0][0] if len(heap) == k else np.inf
        if diff ** 2 < worst:                      # can the far side still contain a closer point?
            search(far)

    search(root)
    return sorted((np.sqrt(-d2), i) for d2, i in heap)

names = "ABCDEFGHIJK"
points = np.array([(6, 2, 5), (4, 9, 9), (1, 4, 6), (2, 8, 3), (5, 11, 11), (3, 10, 10),
                   (9, 5, 4), (8, 1, 2), (7, 3, 1), (11, 7, 8), (10, 6, 7)], dtype=float)
root = build(points, np.arange(len(points)))
for dist, i in knn(points, root, np.array([0.0, 8.0, 12.0]), k=3):
    print(names[i], points[i], round(dist, 3))
# F [ 3. 10. 10.] 4.123
# B [4. 9. 9.] 5.099
# E [ 5. 11. 11.] 5.916

Python’s heapq is a min-heap, so the distances are stored negated to get a max-heap. Production libraries add a few refinements on top of this:

• Leaf buckets. Instead of splitting down to single points, stop at leaves of 10 to 50 points and scan them linearly. That gives fewer pointers, better cache use, and SIMD-friendly loops.
• Squared distances. Square roots are never needed during the search, only at the very end, if at all.
• Better bounds. The plane test only uses one axis. Libraries such as nanoflann and SciPy track the distance from the query to the node’s whole bounding box incrementally, which prunes more.
• Approximate search. Pruning when $(p_a - q_a)^2 \ge d_{\text{best}}^2 / (1 + \varepsilon)^2$ returns a point at most $(1 + \varepsilon)$ times farther than the true neighbor, often much faster. This is the eps parameter in SciPy and FLANN.

How fast, and when it stops working

For a feel of the numbers, with SciPy’s cKDTree on one CPU core and a million random 3D points in a $100 \times 100 \times 100$ m cube:

• building the tree takes 0.29 s,
• a nearest-neighbor query takes 1.6 µs (3.6 µs for 10 neighbors),
• brute force takes 16.7 ms per query, about 10,000 times slower.

The curse of dimensionality. The pruning test only works when the search sphere crosses few splitting planes. In high dimensions, almost everything is about equally far away, the sphere crosses nearly every plane, and the tree ends up visiting almost every node. Counting the nodes the exact 1-NN search visits, for 20,000 uniform random points:

A rule of thumb: a k-d tree only beats brute force when $N \gg 2^k$. For 3D points it is ideal. For a 128-dimensional image descriptor, an exact k-d tree is no better than brute force, and you need approximate methods (randomized k-d trees, hierarchical k-means, graph-based indexes like HNSW).

Static by design. A median-balanced tree is built once. Inserting points one by one makes it unbalanced, and deleting points is awkward. That is fine for a fixed map, but a problem for a map that grows with every scan, which brings us to robotics.

How it is used in robotics for KNN search

Robots perceive the world as point clouds, from LiDARs, depth cameras and stereo. Those clouds are 3D, the regime where k-d trees shine, and they are large: a 128-beam LiDAR with 1024 points per ring at 10 Hz produces about 1.3 million points per second. Almost every processing step needs neighbors.

ICP: finding correspondences

Iterative Closest Point aligns a new scan (source) to a map or a previous scan (target). Each iteration first matches every source point to its nearest target point, then solves for the rigid transform that best aligns the matches (see the Gauss-Newton and ICP post). The matching step is a nearest-neighbor query per source point, every iteration, so it usually dominates the run time.

The pattern is always the same: build the tree once on the target, and query it with the moving source. The target does not move during the alignment, so its tree is reused across all iterations, and only the queries change.

from scipy.spatial import cKDTree

tree = cKDTree(target)                                   # once per target
for _ in range(max_iters):
    p = source @ R.T + t
    dist, idx = tree.query(p, distance_upper_bound=0.5)  # 50 cm gate
    valid = np.isfinite(dist)                            # no target point within 0.5 m
    # ... solve for R, t from p[valid] and target[idx[valid]] ...

The distance_upper_bound is more than an outlier gate: it starts the search with $d_{\text{best}} = 0.5$ instead of infinity, so the tree prunes aggressively from the very first node, and queries become faster.

Normal estimation and local geometry

Point-to-plane ICP, ground segmentation and surface reconstruction all need a surface normal at every point. The standard recipe is PCA on the neighborhood (see the SVD and PCA post): find the $k$ nearest neighbors, compute their covariance matrix, and take the eigenvector with the smallest eigenvalue. Locally the surface is a plane, so the direction of least spread is perpendicular to it.

import numpy as np
from scipy.spatial import cKDTree

rng = np.random.default_rng(0)
cloud = rng.standard_normal((20000, 3))
cloud /= np.linalg.norm(cloud, axis=1, keepdims=True)   # points on a 1 m sphere
cloud += 0.005 * rng.standard_normal(cloud.shape)       # 5 mm range noise

tree = cKDTree(cloud)
_, idx = tree.query(cloud, k=20)                         # 20 neighbors for every point
nbrs = cloud[idx]                                        # (N, 20, 3)
centered = nbrs - nbrs.mean(axis=1, keepdims=True)
cov = np.einsum('nki,nkj->nij', centered, centered) / 20
eigval, eigvec = np.linalg.eigh(cov)                     # eigenvalues in ascending order
normals = eigvec[:, :, 0]                                # direction of least variance

On this noisy sphere, the median normal error is about 2.7°. The same eigenvalues are useful too: $\lambda_0 / (\lambda_0 + \lambda_1 + \lambda_2)$ is near zero on flat surfaces and larger on edges and clutter, which is a cheap way to classify points as planar or edge-like. A normal found by PCA has an arbitrary sign. Flip it to point toward the sensor, which is the usual convention.

k-NN or radius search? LiDAR density drops with the square of the range. A fixed $k$ adapts to it: close to the sensor 20 neighbors span a few centimeters, far away they span a meter, so far-away normals get smoothed across corners. A fixed radius gives a consistent physical scale, but returns too few points far away and huge sets close by. Many pipelines use a hybrid: at most $k$ neighbors within radius $r$ (Open3D’s KDTreeSearchParamHybrid).

Filtering: removing outliers

Dust, rain, and mixed pixels at depth edges create isolated points floating in space. Statistical outlier removal computes the mean distance from each point to its $k$ nearest neighbors, then drops points whose mean distance is much larger than average:

cloud = np.vstack([cloud, rng.uniform(-1.5, 1.5, (200, 3))])   # add 200 outliers

tree = cKDTree(cloud)
d, _ = tree.query(cloud, k=9)                 # k=9 because the first neighbor is the point itself
mean_d = d[:, 1:].mean(axis=1)
keep = mean_d < mean_d.mean() + 2 * mean_d.std()

This removes 171 points, all of them actual outliers. The 29 outliers that survive happen to sit close to the sphere’s surface, where they cannot be told apart from noise by density alone. Radius outlier removal is the radius-search version: drop points that have fewer than $m$ neighbors within radius $r$.

Clustering: segmenting obstacles

After removing the ground, a robot wants to separate the remaining points into objects: a person, a pallet, a pole. Euclidean cluster extraction is a flood fill. Start from an unvisited point, add every point within radius $r$ of it using a radius search, then repeat from the newly added points until the cluster stops growing. Each point triggers one radius query, so the k-d tree makes the whole segmentation roughly $O(N \log N)$ instead of $O(N^2)$. This is pcl::EuclideanClusterExtraction in PCL, and DBSCAN is the same idea with a density threshold.

Motion planning: the nearest node in a growing tree

Sampling-based planners like RRT add one random sample per iteration and connect it to the nearest existing node in the tree. RRT* and PRM also need all nodes within a radius to rewire or connect. Here the points live in configuration space, for example 7D for a robot arm, and three things change:

• The tree grows constantly. A static k-d tree has to be rebuilt periodically, or replaced by a structure that supports insertion. OMPL’s default is GNAT, a metric tree that supports insertion, and it can optionally use FLANN’s k-d trees.
• The metric is not always Euclidean. Joint angles wrap around at $\pm\pi$, and orientations live on $SO(3)$. The plain axis-split k-d tree assumes a flat space, so these coordinates need special handling or a metric tree.
• The dimension is higher. A 7D configuration space still benefits from a k-d tree, but the pruning is much weaker than in 3D, as the dimension table above shows.

Feature matching and place recognition

Visual SLAM matches keypoint descriptors between frames, and loop closure compares a whole scan descriptor against every place visited before. These descriptors have 32 to 256 dimensions, far beyond where exact k-d trees help. Libraries like FLANN use randomized k-d trees: several trees, each splitting on randomly chosen high-variance axes, searched in parallel with a fixed budget of leaf checks. The answer is approximate, but the right match is found most of the time, which is good enough once RANSAC filters the rest.

Keeping up with a growing map

LiDAR SLAM adds a new scan to its map at 10 Hz or more, and rebuilding the tree from scratch is expensive: 0.29 s for a million points, longer than the time between two scans. Systems handle this in a few ways:

• Local maps. Only keep the map near the robot, for example within 100 m, and rebuild a small tree periodically or in a background thread.
• Incremental k-d trees. FAST-LIO2’s ikd-Tree supports point insertion and box-wise deletion, and rebalances only the subtrees that become too lopsided. It can also downsample on insertion, so the map does not grow without bound.
• Voxel hash maps. KISS-ICP and Faster-LIO (iVox) drop the tree entirely. They hash points into voxels and only search the voxels around the query. Insertion and deletion are $O(1)$, and the search is approximate at voxel boundaries, which is perfectly fine for ICP.

Libraries

Practical pitfalls

• Build once, query many. Rebuilding the target tree inside an ICP loop is a classic slowdown. The tree only needs rebuilding when the target changes.
• Batch your queries. One call with all query points lets libraries parallelize and keep the tree hot in cache. In Python, a per-point loop is dominated by interpreter overhead.
• Mind the units. The tree uses one Euclidean metric for all axes. If you mix meters with radians or with intensity values, scale them first, or the neighbors will be meaningless.
• Downsample first. A voxel grid filter before building the tree reduces $N$ and evens out the density, which helps both speed and normal quality.
• Gate the search. A maximum distance (distance_upper_bound, max_distance) speeds up the search and rejects bad correspondences at the same time.

The k-d tree is a simple data structure built on one geometric observation: if a splitting plane is farther away than your best candidate, nothing behind it can do better. In 3D, that observation turns a linear scan into a logarithmic search, and that is a big part of why LiDAR odometry, mapping and perception run in real time.