From a4a75059f302de2a00971f1f485fcf4389710628 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sun, 19 Apr 2020 16:55:17 -0400 Subject: metric/forest: Implement dynamized forests --- src/metric/forest.rs | 152 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 152 insertions(+) create mode 100644 src/metric/forest.rs (limited to 'src/metric') diff --git a/src/metric/forest.rs b/src/metric/forest.rs new file mode 100644 index 0000000..f23c451 --- /dev/null +++ b/src/metric/forest.rs @@ -0,0 +1,152 @@ +//! [Dynamization](https://en.wikipedia.org/wiki/Dynamization) for nearest neighbor search. + +use super::kd::KdTree; +use super::vp::VpTree; +use super::{Metric, NearestNeighbors, Neighborhood}; + +use std::iter::{Extend, Flatten, FromIterator}; + +/// A dynamic wrapper for a static nearest neighbor search data structure. +/// +/// This type applies [dynamization](https://en.wikipedia.org/wiki/Dynamization) to an arbitrary +/// nearest neighbor search structure `T`, allowing new items to be added dynamically. +#[derive(Debug)] +pub struct Forest(Vec>); + +impl Forest +where + U: FromIterator + IntoIterator, +{ + /// Create a new empty forest. + pub fn new() -> Self { + Self(Vec::new()) + } + + /// Add a new item to the forest. + pub fn push(&mut self, item: T) { + let mut items = vec![item]; + + for slot in &mut self.0 { + match slot.take() { + // Collect the items from any trees we encounter... + Some(tree) => { + items.extend(tree); + } + // ... and put them all in the first empty slot + None => { + *slot = Some(items.into_iter().collect()); + return; + } + } + } + + self.0.push(Some(items.into_iter().collect())); + } + + /// Get the number of items in the forest. + pub fn len(&self) -> usize { + let mut len = 0; + for (i, slot) in self.0.iter().enumerate() { + if slot.is_some() { + len |= 1 << i; + } + } + len + } +} + +impl Extend for Forest +where + U: FromIterator + IntoIterator, +{ + fn extend>(&mut self, items: I) { + for item in items { + self.push(item); + } + } +} + +impl FromIterator for Forest +where + U: FromIterator + IntoIterator, +{ + fn from_iter>(items: I) -> Self { + let mut forest = Self::new(); + forest.extend(items); + forest + } +} + +type IntoIterImpl = Flatten>>>; + +/// An iterator that moves items out of a forest. +pub struct IntoIter(IntoIterImpl); + +impl Iterator for IntoIter { + type Item = T::Item; + + fn next(&mut self) -> Option { + self.0.next() + } +} + +impl IntoIterator for Forest { + type Item = T::Item; + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + IntoIter(self.0.into_iter().flatten().flatten()) + } +} + +impl NearestNeighbors for Forest +where + U: Metric, + V: NearestNeighbors, +{ + fn search<'a, 'b, N>(&'a self, neighborhood: N) -> N + where + T: 'a, + U: 'b, + N: Neighborhood<&'a T, &'b U>, + { + self.0 + .iter() + .flatten() + .fold(neighborhood, |n, t| t.search(n)) + } +} + +/// A forest of k-d trees. +pub type KdForest = Forest>; + +/// A forest of vantage-point trees. +pub type VpForest = Forest>; + +#[cfg(test)] +mod tests { + use super::*; + + use crate::metric::tests::test_nearest_neighbors; + use crate::metric::ExhaustiveSearch; + + #[test] + fn test_exhaustive_forest() { + test_nearest_neighbors(Forest::>::from_iter); + } + + #[test] + fn test_forest_forest() { + test_nearest_neighbors(Forest::>>::from_iter); + } + + #[test] + fn test_kd_forest() { + test_nearest_neighbors(KdForest::from_iter); + } + + #[test] + fn test_vp_forest() { + test_nearest_neighbors(VpForest::from_iter); + } +} -- cgit v1.2.3