Cleaner, min leaf size

This commit is contained in:
Thomas Forgione 2018-03-05 11:44:43 +01:00
parent 4c773c290d
commit 64f25f51bd
No known key found for this signature in database
GPG Key ID: C75CD416BD1FFCE1
2 changed files with 35 additions and 27 deletions

View File

@ -50,7 +50,7 @@ fn main() {
Vector2::new(7.0, 2.0),
];
let kdtree = KdTree::new(&mut points, 2);
let kdtree = KdTree::new(&mut points, 2, 1);
println!("{}", kdtree);

View File

@ -15,18 +15,17 @@ impl<T: num::Zero> Zero for T {
}
/// A k-dimension kd tree
pub struct KdTree<T> {
location: T,
left: Option<Box<KdTree<T>>>,
right: Option<Box<KdTree<T>>>,
pub enum KdTree<T> {
Node { location: T, left: Box<KdTree<T>>, right: Box<KdTree<T>> },
Leaf(Vec<T>),
}
impl<T, U> KdTree<T> where T: Copy + Clone + Index<usize, Output = U> + Zero, U: PartialOrd {
pub fn new(points: &mut [T], dimension: usize) -> KdTree<T> {
KdTree::new_with_depth(points, dimension, 0)
pub fn new(points: &mut [T], dimension: usize, leaf_size: usize) -> KdTree<T> {
KdTree::new_with_depth(points, dimension, leaf_size, 0)
}
fn new_with_depth(points: &mut [T], dimension: usize, depth: usize) -> KdTree<T> {
fn new_with_depth(points: &mut [T], dimension: usize, leaf_size: usize, depth: usize) -> KdTree<T> {
// Select axis to work on
let axis = depth % dimension;
@ -36,41 +35,50 @@ impl<T, U> KdTree<T> where T: Copy + Clone + Index<usize, Output = U> + Zero, U:
// Find the median
let len = points.len();
// If there are not enough elements for a leaf, make a leaf
if len <= leaf_size {
let mut contained = vec![Zero::zero(); len];
contained.copy_from_slice(points);
return KdTree::Leaf(contained);
}
// Else split what's remaining
let delimiter = points.len() / 2;
let median = points[delimiter];
// Build the left and right branches
let zero = Zero::zero();
let mut left = vec![zero; delimiter];
left.copy_from_slice(&points[0..delimiter]);
let right = &mut points[delimiter + 1 .. len];
if left.len() == 0 || right.len() == 0 {
KdTree {
// Recursive call
KdTree::Node {
location: median,
left: None,
right: None,
}
} else {
KdTree {
location: median,
left: Some(Box::new(KdTree::new_with_depth(&mut left, dimension, depth + 1))),
right: Some(Box::new(KdTree::new_with_depth(right, dimension, depth + 1))),
}
left: Box::new(KdTree::new_with_depth(&mut left, dimension, leaf_size, depth + 1)),
right: Box::new(KdTree::new_with_depth(right, dimension, leaf_size, depth + 1)),
}
}
}
impl<T: Zero + Copy + Clone + fmt::Display + Index<usize>> KdTree<T> {
fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result {
writeln!(formatter, "{}{}", indent, self.location)?;
if let Some(left) = self.left.as_ref() {
match self {
&KdTree::Node { location, ref left, ref right } => {
writeln!(formatter, "{}{}", indent, location)?;
left.print(formatter, &format!("{} ", indent))?;
}
if let Some(right) = self.right.as_ref() {
right.print(formatter, &format!("{} ", indent))?;
},
&KdTree::Leaf(ref elements) => {
write!(formatter, "{}[", indent)?;
for element in elements {
write!(formatter, "{}, ", element)?;
}
writeln!(formatter, "]")?;
},
}
Ok(())