If you find any mistakes here or anywhere else on my website, email website@CoryHufford.com
This is a simple reference K-D tree implementation in Rust. It should perform reasonably well for smallish trees and queries. The lack of handling for equidistant points in this implementation means its results are, in general, not reproducible.
#[must_use]
fn squared_euclidean<const K: usize>(a: &[f64; K], b: &[f64; K]) -> f64 {
a.iter().zip(b).map(|(ai, bi)| (ai - bi).powi(2)).sum::<f64>()
}
mod kd {
pub(super) const FIRST_AXIS: usize = 0;
const SMALL_TREE_CUTOFF: usize = 8;
pub(super) struct HeapEntry<'a, const K: usize>(
pub(super) &'a [f64; K],
pub(super) f64,
);
impl<const K: usize> PartialEq for HeapEntry<'_, K> {
fn eq(&self, other: &Self) -> bool {
self.1.total_cmp(&other.1).is_eq()
}
}
impl<const K: usize> Eq for HeapEntry<'_, K>{}
impl<const K: usize> PartialOrd for HeapEntry<'_, K> {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl<const K: usize> Ord for HeapEntry<'_, K> {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.1.total_cmp(&other.1)
}
}
pub(super) fn partition<const K: usize>(points: &mut [[f64; K]], axis: usize) {
if points.len() <= SMALL_TREE_CUTOFF {
return;
}
let i = points.len() / 2;
points.select_nth_unstable_by(i, |a, b| a[axis].total_cmp(&b[axis]));
let next_axis = (axis + 1) % K;
partition(&mut points[..i], next_axis);
partition(&mut points[i + 1..], next_axis);
}
pub(super) fn find_nearest_one<'a, const K: usize>(
points: &'a [[f64; K]],
query: &[f64; K],
nearest: &mut (&'a [f64; K], f64),
axis: usize,
) {
if points.len() <= SMALL_TREE_CUTOFF {
for n in points {
let d = super::squared_euclidean(query, n);
if d < nearest.1 {
*nearest = (n, d);
}
}
return;
}
let i = points.len() / 2;
let distance = super::squared_euclidean(query, &points[i]);
if distance < nearest.1 {
*nearest = (&points[i], distance);
}
let split_val = points[i][axis];
let (side_a, side_b) = if query[axis] < split_val {
(&points[..i], &points[i + 1..])
} else {
(&points[i + 1..], &points[..i])
};
let next_axis = (axis + 1) % K;
find_nearest_one(side_a, query, nearest, next_axis);
if (query[axis] - split_val).powi(2) < nearest.1 {
find_nearest_one(side_b, query, nearest, next_axis);
}
}
pub(super) fn find_nearest_n_unsorted<'a, const K: usize>(
points: &'a [[f64; K]],
query: &[f64; K],
n: usize,
nearest_n: &mut std::collections::BinaryHeap<HeapEntry<'a, K>>,
axis: usize,
) {
debug_assert!(n != 0);
if points.len() <= SMALL_TREE_CUTOFF {
for p in points {
let d = super::squared_euclidean(query, p);
if nearest_n.len() < n || d < nearest_n.peek().unwrap().1 {
if nearest_n.len() == n {
nearest_n.pop();
}
nearest_n.push(HeapEntry(p, d));
}
}
return;
}
let i = points.len() / 2;
let distance = super::squared_euclidean(query, &points[i]);
if nearest_n.len() < n || distance < nearest_n.peek().unwrap().1 {
if nearest_n.len() == n {
nearest_n.pop();
}
nearest_n.push(HeapEntry(&points[i], distance));
}
let split_val = points[i][axis];
let (side_a, side_b) = if query[axis] < split_val {
(&points[..i], &points[i + 1..])
} else {
(&points[i + 1..], &points[..i])
};
let next_axis = (axis + 1) % K;
find_nearest_n_unsorted(side_a, query, n, nearest_n, next_axis);
if nearest_n.len() < n
|| (query[axis] - split_val).powi(2) < nearest_n.peek().unwrap().1
{
find_nearest_n_unsorted(side_b, query, n, nearest_n, next_axis);
}
}
pub(super) fn find_within_distance_squared_unsorted<const K: usize>(
points: &[[f64; K]],
query: &[f64; K],
distance: f64,
distance_squared: f64,
results: &mut Vec<([f64; K], f64)>,
axis: usize,
) {
if points.len() <= SMALL_TREE_CUTOFF {
for n in points {
let d = super::squared_euclidean(query, n);
if d <= distance_squared {
results.push((*n, d));
}
}
return;
}
let i = points.len() / 2;
let next_axis = (axis + 1) % K;
let lower_bound = query[axis] - distance;
let upper_bound = query[axis] + distance;
let split_val = points[i][axis];
let d = super::squared_euclidean(query, &points[i]);
if d <= distance_squared {
results.push((points[i], d));
}
if split_val >= lower_bound {
find_within_distance_squared_unsorted(
&points[..i],
query,
distance,
distance_squared,
results,
next_axis,
);
}
if split_val <= upper_bound {
find_within_distance_squared_unsorted(
&points[i + 1..],
query,
distance,
distance_squared,
results,
next_axis,
);
}
}
}
/// 2-D Euclidean distance K-D tree.
pub struct KdTree {
points: Vec<[f64; 2]>,
}
impl KdTree {
/// Can be empty.
#[must_use]
pub fn new(mut points: Vec<[f64; 2]>) -> Self {
debug_assert!(points.iter().all(|a| a.iter().all(|f| f.is_finite())));
kd::partition(&mut points, kd::FIRST_AXIS);
Self{points}
}
#[must_use]
pub fn len(&self) -> usize {
self.points.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.points.is_empty()
}
/// Returns the nearest point to `query` in the tree and its distance from
/// `query`, or nothing if the tree is empty. If there are multiple points
/// equidistant from `query`, it's undefined which one is returned.
#[must_use]
pub fn nearest_one(&self, query: &[f64; 2]) -> Option<([f64; 2], f64)> {
debug_assert!(query.iter().all(|f| f.is_finite()));
if self.points.is_empty() {
return None;
}
let mut nearest = (&self.points[0], f64::INFINITY);
kd::find_nearest_one(&self.points, query, &mut nearest, kd::FIRST_AXIS);
Some((*nearest.0, nearest.1.sqrt()))
}
/// Returns the `n` nearest points to `query`, along with their distances
/// from `query`, in no particular order. Fewer than `n` points will be
/// returned if the tree contains fewer than `n` points. If there are many
/// points equidistant from `query`, it's undefined which ones are returned.
#[must_use]
pub fn nearest_n_unsorted(
&self,
query: &[f64; 2],
n: usize,
) -> Vec<([f64; 2], f64)> {
debug_assert!(query.iter().all(|f| f.is_finite()));
let n = n.min(self.points.len());
if n == 0 {
return Vec::new();
}
let mut nearest_n = std::collections::BinaryHeap::<kd::HeapEntry<2>>::with_capacity(n);
kd::find_nearest_n_unsorted(&self.points, query, n, &mut nearest_n, kd::FIRST_AXIS);
nearest_n.into_iter().map(|he| (*he.0, he.1.sqrt())).collect()
}
/// Returns the `n` nearest points to `query`, along with their distances
/// from `query`, in order from closest to `query` to farthest. Fewer than
/// `n` points will be returned if the tree contains fewer than `n` points.
/// If there are many points equidistant from `query`, it's undefined which
/// ones are returned.
#[must_use]
pub fn nearest_n_sorted(
&self,
query: &[f64; 2],
n: usize,
) -> Vec<([f64; 2], f64)> {
debug_assert!(query.iter().all(|f| f.is_finite()));
let mut result = self.nearest_n_unsorted(query, n);
result.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
result
}
/// Returns all points within `distance` of `query`, inclusive, along with
/// their distances from `query`, in no particular order.
#[must_use]
pub fn within_distance_unsorted(
&self,
query: &[f64; 2],
distance: f64,
) -> Vec<([f64; 2], f64)> {
debug_assert!(query.iter().all(|f| f.is_finite()));
debug_assert!(distance.is_finite());
if distance < 0.0 {
return Vec::new();
}
let mut results = Vec::new();
kd::find_within_distance_squared_unsorted(
&self.points,
query,
distance,
distance.powi(2),
&mut results,
kd::FIRST_AXIS,
);
for (_, d) in &mut results {
*d = d.sqrt();
}
results
}
/// Returns all points within `distance` of `query`, inclusive, along with
/// their distances from `query`, in order from closest to `query` to
/// farthest.
#[must_use]
pub fn within_distance_sorted(
&self,
query: &[f64; 2],
distance: f64,
) -> Vec<([f64; 2], f64)> {
debug_assert!(query.iter().all(|f| f.is_finite()));
debug_assert!(distance.is_finite());
let mut points = self.within_distance_unsorted(query, distance);
points.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
points
}
}