summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--src/vp.rs92
1 files changed, 92 insertions, 0 deletions
diff --git a/src/vp.rs b/src/vp.rs
index 67a094f..bf91d57 100644
--- a/src/vp.rs
+++ b/src/vp.rs
@@ -201,6 +201,11 @@ impl<T: Proximity> VpTree<T> {
}
}
+ /// Iterate over the items stored in this tree.
+ pub fn iter(&self) -> Iter<T> {
+ (&self).into_iter()
+ }
+
/// Rebalance this VP tree.
pub fn balance(&mut self) {
let mut nodes = Vec::new();
@@ -327,6 +332,56 @@ impl<T: Proximity> IntoIterator for VpTree<T> {
}
}
+/// An iterator over the values in a VP tree.
+pub struct Iter<'a, T: Proximity> {
+ stack: Vec<&'a VpNode<T>>,
+}
+
+impl<'a, T: Proximity> Iter<'a, T> {
+ fn new(node: &'a Option<VpNode<T>>) -> Self {
+ Self {
+ stack: node.as_ref().into_iter().collect(),
+ }
+ }
+}
+
+impl<'a, T> Debug for Iter<'a, T>
+where
+ T: Proximity + Debug,
+ DistanceValue<T>: Debug,
+{
+ fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+ f.debug_struct("Iter")
+ .field("stack", &self.stack)
+ .finish()
+ }
+}
+
+impl<'a, T: Proximity> Iterator for Iter<'a, T> {
+ type Item = &'a T;
+
+ fn next(&mut self) -> Option<&'a T> {
+ self.stack.pop().map(|node| {
+ if let Some(inside) = &node.inside {
+ self.stack.push(inside);
+ }
+ if let Some(outside) = &node.outside {
+ self.stack.push(outside);
+ }
+ &node.item
+ })
+ }
+}
+
+impl<'a, T: Proximity> IntoIterator for &'a VpTree<T> {
+ type Item = &'a T;
+ type IntoIter = Iter<'a, T>;
+
+ fn into_iter(self) -> Self::IntoIter {
+ Iter::new(&self.root)
+ }
+}
+
impl<K, V> NearestNeighbors<K, V> for VpTree<V>
where
K: Proximity<V, Distance = V::Distance>,
@@ -452,6 +507,11 @@ impl<T: Proximity> FlatVpTree<T> {
nodes: FlatVpNode::balanced(items),
}
}
+
+ /// Iterate over the items stored in this tree.
+ pub fn iter(&self) -> FlatIter<T> {
+ (&self).into_iter()
+ }
}
impl<T> Debug for FlatVpTree<T>
@@ -504,6 +564,38 @@ impl<T: Proximity> IntoIterator for FlatVpTree<T> {
}
}
+/// An iterator over the values in a flat VP tree.
+pub struct FlatIter<'a, T: Proximity>(std::slice::Iter<'a, FlatVpNode<T>>);
+
+impl<'a, T> Debug for FlatIter<'a, T>
+where
+ T: Proximity + Debug,
+ DistanceValue<T>: Debug,
+{
+ fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+ f.debug_tuple("FlatIter")
+ .field(&self.0)
+ .finish()
+ }
+}
+
+impl<'a, T: Proximity> Iterator for FlatIter<'a, T> {
+ type Item = &'a T;
+
+ fn next(&mut self) -> Option<&'a T> {
+ self.0.next().map(|n| &n.item)
+ }
+}
+
+impl<'a, T: Proximity> IntoIterator for &'a FlatVpTree<T> {
+ type Item = &'a T;
+ type IntoIter = FlatIter<'a, T>;
+
+ fn into_iter(self) -> Self::IntoIter {
+ FlatIter(self.nodes.iter())
+ }
+}
+
impl<K, V> NearestNeighbors<K, V> for FlatVpTree<V>
where
K: Proximity<V, Distance = V::Distance>,