Give info about the tree

This commit is contained in:
Thomas Forgione 2018-03-08 16:33:03 +01:00
parent a4176bca5a
commit 5f78e5fc5f
No known key found for this signature in database
GPG Key ID: C75CD416BD1FFCE1
1 changed files with 34 additions and 5 deletions

View File

@ -1,6 +1,11 @@
use std::fmt; use std::fmt;
use std::cmp::Ordering; use std::cmp::Ordering;
pub enum Side {
Left,
Right,
}
/// A k-dimension kd tree /// A k-dimension kd tree
pub enum KdTree<T> { pub enum KdTree<T> {
Node { Node {
@ -18,6 +23,7 @@ impl<T: Clone> KdTree<T> where {
comparator: F, comparator: F,
leaf_size: usize, leaf_size: usize,
) -> KdTree<T> { ) -> KdTree<T> {
// Compute bounding box of points
KdTree::new_with_depth(points, dimension, &comparator, leaf_size, 0) KdTree::new_with_depth(points, dimension, &comparator, leaf_size, 0)
} }
@ -77,16 +83,39 @@ impl<T: Clone> KdTree<T> where {
)), )),
} }
} }
// Fn(location, elements)
pub fn traverse_leaves<F: FnMut(&Vec<(T, Side)>, &Vec<T>)>(&self, callback: &mut F) {
self.traverse_leaves_aux(callback, &mut vec![])
}
pub fn traverse_leaves_aux<F:FnMut(&Vec<(T, Side)>, &Vec<T>)>(&self, callback: &mut F, locations: &mut Vec<(T, Side)>) {
// Compute the full bounding box of the points
match *self {
KdTree::Node { ref location , ref left, ref right } => {
locations.push((location.clone(), Side::Left));
left.traverse_leaves_aux(callback, locations);
locations.pop();
locations.push((location.clone(), Side::Right));
right.traverse_leaves_aux(callback, locations);
locations.pop();
},
KdTree::Leaf(ref elements) => {
callback(locations, elements);
}
}
}
} }
impl<T: fmt::Display> KdTree<T> { impl<T: fmt::Display> KdTree<T> {
fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result { fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result {
match self { match self {
&KdTree::Node { &KdTree::Node { ref location, ref left, ref right, } => {
ref location,
ref left,
ref right,
} => {
writeln!(formatter, "{}{}", indent, location)?; writeln!(formatter, "{}{}", indent, location)?;
left.print(formatter, &format!("{} ", indent))?; left.print(formatter, &format!("{} ", indent))?;
right.print(formatter, &format!("{} ", indent))?; right.print(formatter, &format!("{} ", indent))?;