summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--benches/benches.rs8
-rw-r--r--src/kd.rs155
2 files changed, 162 insertions, 1 deletions
diff --git a/benches/benches.rs b/benches/benches.rs
index 8464179..3e67c2a 100644
--- a/benches/benches.rs
+++ b/benches/benches.rs
@@ -2,7 +2,7 @@
use acap::euclid::Euclidean;
use acap::exhaustive::ExhaustiveSearch;
-use acap::kd::KdTree;
+use acap::kd::{FlatKdTree, KdTree};
use acap::vp::{FlatVpTree, VpTree};
use acap::NearestNeighbors;
@@ -39,6 +39,7 @@ fn bench_from_iter(c: &mut Criterion) {
group.bench_function("VpTree", |b| b.iter(|| VpTree::from_iter(points.clone())));
group.bench_function("FlatVpTree", |b| b.iter(|| FlatVpTree::from_iter(points.clone())));
group.bench_function("KdTree", |b| b.iter(|| KdTree::from_iter(points.clone())));
+ group.bench_function("FlatKdTree", |b| b.iter(|| FlatKdTree::from_iter(points.clone())));
group.finish();
}
@@ -50,12 +51,14 @@ fn bench_nearest_neighbors(c: &mut Criterion) {
let vp_tree = VpTree::from_iter(points.clone());
let flat_vp_tree = FlatVpTree::from_iter(points.clone());
let kd_tree = KdTree::from_iter(points.clone());
+ let flat_kd_tree = FlatKdTree::from_iter(points.clone());
let mut nearest = c.benchmark_group("NearestNeighbors::nearest");
nearest.bench_function("ExhaustiveSearch", |b| b.iter(|| exhaustive.nearest(&target)));
nearest.bench_function("VpTree", |b| b.iter(|| vp_tree.nearest(&target)));
nearest.bench_function("FlatVpTree", |b| b.iter(|| flat_vp_tree.nearest(&target)));
nearest.bench_function("KdTree", |b| b.iter(|| kd_tree.nearest(&target)));
+ nearest.bench_function("FlatKdTree", |b| b.iter(|| flat_kd_tree.nearest(&target)));
nearest.finish();
let mut nearest_within = c.benchmark_group("NearestNeighbors::nearest_within");
@@ -63,6 +66,7 @@ fn bench_nearest_neighbors(c: &mut Criterion) {
nearest_within.bench_function("VpTree", |b| b.iter(|| vp_tree.nearest_within(&target, 0.1)));
nearest_within.bench_function("FlatVpTree", |b| b.iter(|| flat_vp_tree.nearest_within(&target, 0.1)));
nearest_within.bench_function("KdTree", |b| b.iter(|| kd_tree.nearest_within(&target, 0.1)));
+ nearest_within.bench_function("FlatKdTree", |b| b.iter(|| flat_kd_tree.nearest_within(&target, 0.1)));
nearest_within.finish();
let mut k_nearest = c.benchmark_group("NearestNeighbors::k_nearest");
@@ -70,6 +74,7 @@ fn bench_nearest_neighbors(c: &mut Criterion) {
k_nearest.bench_function("VpTree", |b| b.iter(|| vp_tree.k_nearest(&target, 3)));
k_nearest.bench_function("FlatVpTree", |b| b.iter(|| flat_vp_tree.k_nearest(&target, 3)));
k_nearest.bench_function("KdTree", |b| b.iter(|| kd_tree.k_nearest(&target, 3)));
+ k_nearest.bench_function("FlatKdTree", |b| b.iter(|| flat_kd_tree.k_nearest(&target, 3)));
k_nearest.finish();
let mut k_nearest_within = c.benchmark_group("NearestNeighbors::k_nearest_within");
@@ -77,6 +82,7 @@ fn bench_nearest_neighbors(c: &mut Criterion) {
k_nearest_within.bench_function("VpTree", |b| b.iter(|| vp_tree.k_nearest_within(&target, 3, 0.1)));
k_nearest_within.bench_function("FlatVpTree", |b| b.iter(|| flat_vp_tree.k_nearest_within(&target, 3, 0.1)));
k_nearest_within.bench_function("KdTree", |b| b.iter(|| kd_tree.k_nearest_within(&target, 3, 0.1)));
+ k_nearest_within.bench_function("FlatKdTree", |b| b.iter(|| flat_kd_tree.k_nearest_within(&target, 3, 0.1)));
k_nearest_within.finish();
}
diff --git a/src/kd.rs b/src/kd.rs
index 97616e7..4f591f9 100644
--- a/src/kd.rs
+++ b/src/kd.rs
@@ -333,6 +333,156 @@ where
V: Coordinates,
{}
+/// A node in a flat k-d tree.
+#[derive(Debug)]
+struct FlatKdNode<T> {
+ /// The vantage point itself.
+ item: T,
+ /// The size of the left subtree.
+ left_len: usize,
+}
+
+impl<T: Coordinates> FlatKdNode<T> {
+ /// Create a new FlatKdNode.
+ fn new(item: T) -> Self {
+ Self {
+ item,
+ left_len: 0,
+ }
+ }
+
+ /// Create a balanced tree.
+ fn balanced<I: IntoIterator<Item = T>>(items: I) -> Vec<Self> {
+ let mut nodes: Vec<_> = items
+ .into_iter()
+ .map(Self::new)
+ .collect();
+
+ Self::balance_recursive(&mut nodes, 0);
+
+ nodes
+ }
+
+ /// Create a balanced subtree.
+ fn balance_recursive(nodes: &mut [Self], level: usize) {
+ if !nodes.is_empty() {
+ nodes.sort_by_cached_key(|x| Ordered::new(x.item.coord(level)));
+
+ let mid = nodes.len() / 2;
+ nodes.swap(0, mid);
+
+ let (node, children) = nodes.split_first_mut().unwrap();
+ let (left, right) = children.split_at_mut(mid);
+ node.left_len = left.len();
+
+ let next = (level + 1) % node.item.dims();
+ Self::balance_recursive(left, next);
+ Self::balance_recursive(right, next);
+ }
+ }
+}
+
+impl<'a, K, V, N> KdSearch<K, &'a V, N> for &'a [FlatKdNode<V>]
+where
+ K: KdProximity<&'a V>,
+ V: Coordinates,
+ N: Neighborhood<K, &'a V>,
+{
+ fn item(self) -> &'a V {
+ &self[0].item
+ }
+
+ fn left(self) -> Option<Self> {
+ let end = self[0].left_len + 1;
+ if end > 1 {
+ Some(&self[1..end])
+ } else {
+ None
+ }
+ }
+
+ fn right(self) -> Option<Self> {
+ let start = self[0].left_len + 1;
+ if start < self.len() {
+ Some(&self[start..])
+ } else {
+ None
+ }
+ }
+}
+
+/// A [k-d tree] stored as a flat array.
+///
+/// A FlatKdTree is always balanced and usually more efficient than a [KdTree], but doesn't support
+/// dynamic updates.
+///
+/// [k-d tree]: https://en.wikipedia.org/wiki/K-d_tree
+#[derive(Debug)]
+pub struct FlatKdTree<T> {
+ nodes: Vec<FlatKdNode<T>>,
+}
+
+impl<T: Coordinates> FlatKdTree<T> {
+ /// Create a balanced tree out of a sequence of items.
+ pub fn balanced<I: IntoIterator<Item = T>>(items: I) -> Self {
+ Self {
+ nodes: FlatKdNode::balanced(items),
+ }
+ }
+}
+
+impl<T: Coordinates> FromIterator<T> for FlatKdTree<T> {
+ fn from_iter<I: IntoIterator<Item = T>>(items: I) -> Self {
+ Self::balanced(items)
+ }
+}
+
+/// An iterator that moves values out of a flat k-d tree.
+#[derive(Debug)]
+pub struct FlatIntoIter<T>(std::vec::IntoIter<FlatKdNode<T>>);
+
+impl<T> Iterator for FlatIntoIter<T> {
+ type Item = T;
+
+ fn next(&mut self) -> Option<T> {
+ self.0.next().map(|n| n.item)
+ }
+}
+
+impl<T> IntoIterator for FlatKdTree<T> {
+ type Item = T;
+ type IntoIter = FlatIntoIter<T>;
+
+ fn into_iter(self) -> Self::IntoIter {
+ FlatIntoIter(self.nodes.into_iter())
+ }
+}
+
+impl<K, V> NearestNeighbors<K, V> for FlatKdTree<V>
+where
+ K: KdProximity<V>,
+ V: Coordinates,
+{
+ fn search<'k, 'v, N>(&'v self, mut neighborhood: N) -> N
+ where
+ K: 'k,
+ V: 'v,
+ N: Neighborhood<&'k K, &'v V>,
+ {
+ if !self.nodes.is_empty() {
+ let mut closest = neighborhood.target().as_vec();
+ self.nodes.as_slice().search(0, &mut closest, &mut neighborhood);
+ }
+ neighborhood
+ }
+}
+
+impl<K, V> ExactNeighbors<K, V> for FlatKdTree<V>
+where
+ K: KdMetric<V>,
+ V: Coordinates,
+{}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -354,4 +504,9 @@ mod tests {
tree
});
}
+
+ #[test]
+ fn test_flat_kd_tree() {
+ test_nearest_neighbors(FlatKdTree::from_iter);
+ }
}