k-d tree

Deep dive

How it works

In: N points in k dimensions (build); a query point and k or a radius, or a box · Out: the k nearest points, points within a radius, or points in the box · State: the points reordered in one array (implicit tree) · Cost: build O(N log N); nearest about 1 µs on 1M uniform 3-D points (measured below)

1. Build on slice [lo, hi) with axis = depth mod k: quickselect the median mid = (lo+hi)/2 on that axis so smaller-or-equal values sit left and greater-or-equal right; recurse on [lo, mid) and [mid+1, hi). Depth is about log2 N (levels = floor(log2 N) + 1).
2. Nearest/k-nearest: visit the node, update a max-heap of the k best; descend first into the side containing the query; visit the other side only if (q[axis] - p[axis])^2 is below the current k-th best squared distance (or the heap is not full).
3. Box query: report the node if inside; go left if lo[axis] <= p[axis], right if hi[axis] >= p[axis].

axis x:            (5,4)
                 /       \
axis y:      (2,3)        (8,1)       slice [lo,hi) = subtree;
            /     \      /     \       median at (lo+hi)/2,
x:        (1,5)  (3,1) (7,6)  (9,2)    no child pointers stored

Code

Implicit tree over [f64; K], k-nearest with a bounded max-heap, box query. Checked against brute force on 20,000 random 3-D points (k = 1, 5, 20, and 200 box queries) and on a duplicate-heavy set.

use std::collections::BinaryHeap;

pub type P<const K: usize> = [f64; K];

/// Implicit k-d tree: the median of each slice [lo, hi) sits at mid = (lo+hi)/2,
/// the split axis is the depth modulo K, children are the slices [lo, mid) and [mid+1, hi).
pub struct KdTree<const K: usize> {
    pts: Vec<P<K>>,
}

fn d2<const K: usize>(a: &P<K>, b: &P<K>) -> f64 {
    (0..K).map(|i| (a[i] - b[i]) * (a[i] - b[i])).sum()
}

fn build<const K: usize>(s: &mut [P<K>], axis: usize) {
    if s.len() <= 1 { return; }
    let mid = s.len() / 2;
    // O(n) expected quickselect: median to s[mid], smaller-or-equal left, greater-or-equal right
    s.select_nth_unstable_by(mid, |a, b| a[axis].total_cmp(&b[axis]));
    let (l, r) = s.split_at_mut(mid);
    let next = (axis + 1) % K;
    build(l, next);
    build(&mut r[1..], next);
}

/// Max-heap entry ordered by distance (squared distances are non-negative, so total_cmp is safe).
struct Hit(f64, usize);
impl PartialEq for Hit { fn eq(&self, o: &Self) -> bool { self.0.total_cmp(&o.0).is_eq() } }
impl Eq for Hit {}
impl PartialOrd for Hit { fn partial_cmp(&self, o: &Self) -> Option<std::cmp::Ordering> { Some(self.cmp(o)) } }
impl Ord for Hit { fn cmp(&self, o: &Self) -> std::cmp::Ordering { self.0.total_cmp(&o.0) } }

impl<const K: usize> KdTree<K> {
    pub fn new(mut pts: Vec<P<K>>) -> Self {
        build(&mut pts, 0);
        KdTree { pts }
    }

    /// Indices (into the internal order) and squared distances of the k nearest points, nearest first.
    pub fn knn(&self, q: &P<K>, k: usize) -> Vec<(usize, f64)> {
        let mut heap = BinaryHeap::new();
        if k > 0 { self.knn_rec(0, self.pts.len(), 0, q, k, &mut heap); }
        let mut v: Vec<_> = heap.into_iter().map(|Hit(d, i)| (i, d)).collect();
        v.sort_by(|a, b| a.1.total_cmp(&b.1));
        v
    }

    fn knn_rec(&self, lo: usize, hi: usize, axis: usize, q: &P<K>, k: usize, heap: &mut BinaryHeap<Hit>) {
        if lo >= hi { return; }
        let mid = (lo + hi) / 2;
        let p = &self.pts[mid];
        let d = d2(p, q);
        if heap.len() < k { heap.push(Hit(d, mid)); }
        else if d < heap.peek().unwrap().0 { heap.pop(); heap.push(Hit(d, mid)); }
        let diff = q[axis] - p[axis];
        let next = (axis + 1) % K;
        let (near, far) = if diff < 0.0 { ((lo, mid), (mid + 1, hi)) } else { ((mid + 1, hi), (lo, mid)) };
        self.knn_rec(near.0, near.1, next, q, k, heap);
        // visit the far side only if the splitting plane is closer than the current k-th best
        if heap.len() < k || diff * diff < heap.peek().unwrap().0 {
            self.knn_rec(far.0, far.1, next, q, k, heap);
        }
    }

    /// All points inside the closed box [lo, hi].
    pub fn range(&self, lo: &P<K>, hi: &P<K>, out: &mut Vec<usize>) {
        self.range_rec(0, self.pts.len(), 0, lo, hi, out);
    }

    fn range_rec(&self, a: usize, b: usize, axis: usize, lo: &P<K>, hi: &P<K>, out: &mut Vec<usize>) {
        if a >= b { return; }
        let mid = (a + b) / 2;
        let p = &self.pts[mid];
        if (0..K).all(|i| lo[i] <= p[i] && p[i] <= hi[i]) { out.push(mid); }
        let next = (axis + 1) % K;
        // duplicates of the split value may lie on both sides, so the tests are non-strict
        if lo[axis] <= p[axis] { self.range_rec(a, mid, next, lo, hi, out); }
        if hi[axis] >= p[axis] { self.range_rec(mid + 1, b, next, lo, hi, out); }
    }

    pub fn point(&self, i: usize) -> &P<K> { &self.pts[i] }
}

Complexity and constants

• Build: each level is O(N) quickselect, so O(N log N); order of 50-230 ms for 1M points, 2-D to 16-D, single thread, f64, safe code (measured, Apple M4 Pro).
• Memory: N * k * 8 B for f64 and nothing else (plus an id or payload per point if you need one).
• Query, 1M uniform random points, single thread, my measurement: 2-D 0.8 µs, 3-D 1.3 µs (nearest), 3-D 2.3 µs (10 nearest); 8-D 49 µs; 16-D 12 ms, which is slower than the 2.9 ms linear scan.
• Cost drivers: nodes visited (each a likely cache miss once the tree exceeds cache) and distance evaluations. Making it cheap means fewer, larger leaves and a contiguous array.
• Box query: O(N^(1-1/k) + T) worst case for T results.

Efficiency

• Cost: build O(N log N), order of 50-230 ms for 1M points in 2-16 dimensions (Apple M4 Pro, single thread, f64, measured above). Nearest on 1M uniform points: 0.8 µs (2-D), 1.3 µs (3-D), 49 µs (8-D); 16-D takes 12 ms against a 2.9 ms linear scan (measured above). Memory N·k·8 B for f64. Bound: memory latency, each visit is a dependent load.
• Levers:
  • bucket leaves of 8-32 points scanned linearly — fewer levels, SIMD-friendly, less memory (defaults: SciPy leafsize=16, nanoflann leaf_max_size=10, scikit-learn leaf_size=30) — always
  • implicit array layout — no child pointers; the array can be saved, mmapped or shared between threads — static data
  • build once, query many, rebuild in batches — build cost amortised over thousands of queries — read-heavy data
  • seed the pruning radius with a known radius — skips most far branches — radius searches
• Naive → standard → state of the art:
  • linear scan · O(N) per query, 2.9 ms at 16-D here, the winner above about 10 dimensions
  • median-split implicit k-d tree with bucket leaves · about 1 µs per query at 3-D
  • a uniform grid or cell index for evenly spread 2-3-D points, or an approximate index when the dimension is high

Implementation notes

• Choose the axis of widest spread (or range) at each node instead of cycling when data is clustered; store the axis per node then. Cycling is fine for uniform data.
• Quickselect leaves equal keys on both sides of the median; use non-strict comparisons in queries (as above) or duplicates are lost.
• Use squared distances throughout; take a square root only for output.
• Use total_cmp, not partial_cmp().unwrap(); one NaN coordinate otherwise panics or corrupts the partition.
• Parallelise by recursing into the two halves with rayon::join or threads below a size threshold; queries are read-only and need no locks.

Pitfalls

• Dimension above about 10 (rule of thumb: N must be much larger than 2^k): nearly every node is visited. My 16-D run was 4x slower than brute force.
• Data on a low-dimensional manifold inside a high-dimensional space still works well; test on your data, not uniform noise.
• Insertion into a balanced tree degrades it; deletion needs tombstones or rebuild.
• Many exact duplicates make median splits useless on that axis; widest-spread axis selection and leaf buckets handle it.
• Geographic data: Euclidean on lon/lat is wrong near the poles and antimeridian; index 3-D unit-sphere coordinates instead and convert the chord distance.

Variants and alternatives

In production

• Scientific computing and ML: SciPy cKDTree/KDTree (docs) and scikit-learn neighbours (docs), for low dimension.
• Robotics and 3-D perception: Point Cloud Library KdTree search (docs); Open3D KDTreeFlann (docs); nanoflann is a header-only C++ library for kd-trees over your own point containers.
• Computational geometry: CGAL spatial searching with kd-trees (docs).
• Graphics: pbrt's third edition has a kd-tree accelerator for ray intersection (book).
• Maps: static point indexes in the browser (kdbush) and Rust (geo-index).

Implementations

Sources

• Jon Louis Bentley, Multidimensional Binary Search Trees Used for Associative Searching, Communications of the ACM 1975. paper
• Jerome H. Friedman, Jon Louis Bentley, Raphael Ari Finkel, An Algorithm for Finding Best Matches in Logarithmic Expected Time, ACM TOMS 1977. paper
• Sunil Arya, David M. Mount, Nathan S. Netanyahu, Ruth Silverman, Angela Y. Wu, An Optimal Algorithm for Approximate Nearest Neighbor Searching in Fixed Dimensions, Journal of the ACM 1998. paper · code
• Marius Muja, David G. Lowe, Fast Approximate Nearest Neighbors with Automatic Algorithm Configuration, VISAPP 2009. paper · code
• SciPy documentation, scipy.spatial.cKDTree. docs
• scikit-learn documentation, Nearest Neighbors. docs