From e272ede323d74d1dd30d28fd862099376b49265b Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Tue, 26 May 2020 12:54:26 -0400 Subject: vp: Implement flat VP trees --- src/vp.rs | 183 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 182 insertions(+), 1 deletion(-) (limited to 'src') diff --git a/src/vp.rs b/src/vp.rs index 120e13b..e0645de 100644 --- a/src/vp.rs +++ b/src/vp.rs @@ -347,6 +347,183 @@ where V: Metric, {} +/// A node in a flat VP tree. +#[derive(Debug)] +struct FlatVpNode> { + /// The vantage point itself. + item: T, + /// The radius of this node. + radius: R, + /// The size of the inside subtree. + inside_len: usize, +} + +impl FlatVpNode { + /// Create a new FlatVpNode. + fn new(item: T) -> Self { + Self { + item, + radius: zero(), + inside_len: 0, + } + } + + /// Create a balanced tree. + fn balanced>(items: I) -> Vec { + let mut nodes: Vec<_> = items + .into_iter() + .map(Self::new) + .collect(); + + Self::balance_recursive(&mut nodes); + + nodes + } + + /// Create a balanced subtree. + fn balance_recursive(nodes: &mut [Self]) { + if let Some((node, children)) = nodes.split_first_mut() { + children.sort_by_cached_key(|x| Ordered::new(node.item.distance(&x.item))); + + let (inside, outside) = children.split_at_mut(children.len() / 2); + if let Some(last) = inside.last() { + node.radius = node.item.distance(&last.item).into(); + } + + node.inside_len = inside.len(); + + Self::balance_recursive(inside); + Self::balance_recursive(outside); + } + } +} + +impl<'a, K, V, N> VpSearch for &'a [FlatVpNode] +where + K: Proximity<&'a V, Distance = V::Distance>, + V: Proximity, + N: Neighborhood, +{ + fn item(self) -> &'a V { + &self[0].item + } + + fn radius(self) -> DistanceValue { + self[0].radius + } + + fn inside(self) -> Option { + let end = self[0].inside_len + 1; + if end > 1 { + Some(&self[1..end]) + } else { + None + } + } + + fn outside(self) -> Option { + let start = self[0].inside_len + 1; + if start < self.len() { + Some(&self[start..]) + } else { + None + } + } +} + +/// A [vantage-point tree] stored as a flat array. +/// +/// A FlatVpTree is always balanced and usually more efficient than a [VpTree], but doesn't support +/// dynamic updates. +/// +/// [vantage-point tree]: https://en.wikipedia.org/wiki/Vantage-point_tree +pub struct FlatVpTree { + nodes: Vec>, +} + +impl FlatVpTree { + /// Create a balanced tree out of a sequence of items. + pub fn balanced>(items: I) -> Self { + Self { + nodes: FlatVpNode::balanced(items), + } + } +} + +impl Debug for FlatVpTree +where + T: Proximity + Debug, + DistanceValue: Debug, +{ + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + f.debug_struct("FlatVpTree") + .field("node", &self.nodes) + .finish() + } +} + +impl FromIterator for FlatVpTree { + fn from_iter>(items: I) -> Self { + Self::balanced(items) + } +} + +/// An iterator that moves values out of a flat VP tree. +pub struct FlatIntoIter(std::vec::IntoIter>); + +impl Debug for FlatIntoIter +where + T: Proximity + Debug, + DistanceValue: Debug, +{ + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + f.debug_tuple("FlatIntoIter") + .field(&self.0) + .finish() + } +} + +impl Iterator for FlatIntoIter { + type Item = T; + + fn next(&mut self) -> Option { + self.0.next().map(|n| n.item) + } +} + +impl IntoIterator for FlatVpTree { + type Item = T; + type IntoIter = FlatIntoIter; + + fn into_iter(self) -> Self::IntoIter { + FlatIntoIter(self.nodes.into_iter()) + } +} + +impl NearestNeighbors for FlatVpTree +where + K: Proximity, + V: Proximity, +{ + 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() { + self.nodes.as_slice().search(&mut neighborhood); + } + neighborhood + } +} + +impl ExactNeighbors for FlatVpTree +where + K: Metric, + V: Metric, +{} + #[cfg(test)] mod tests { use super::*; @@ -368,5 +545,9 @@ mod tests { tree }); } -} + #[test] + fn test_flat_vp_tree() { + test_nearest_neighbors(FlatVpTree::from_iter); + } +} -- cgit v1.2.3