diff --git a/src/example.rs b/src/example.rs index 578398b..5f22feb 100644 --- a/src/example.rs +++ b/src/example.rs @@ -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); diff --git a/src/lib.rs b/src/lib.rs index a7e06bc..bacd90d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,18 +15,17 @@ impl Zero for T { } /// A k-dimension kd tree -pub struct KdTree { - location: T, - left: Option>>, - right: Option>>, +pub enum KdTree { + Node { location: T, left: Box>, right: Box> }, + Leaf(Vec), } impl KdTree where T: Copy + Clone + Index + Zero, U: PartialOrd { - pub fn new(points: &mut [T], dimension: usize) -> KdTree { - KdTree::new_with_depth(points, dimension, 0) + pub fn new(points: &mut [T], dimension: usize, leaf_size: usize) -> KdTree { + KdTree::new_with_depth(points, dimension, leaf_size, 0) } - fn new_with_depth(points: &mut [T], dimension: usize, depth: usize) -> KdTree { + fn new_with_depth(points: &mut [T], dimension: usize, leaf_size: usize, depth: usize) -> KdTree { // Select axis to work on let axis = depth % dimension; @@ -36,41 +35,50 @@ impl KdTree where T: Copy + Clone + Index + 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 { - 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))), - } + // Recursive call + KdTree::Node { + location: median, + 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> KdTree { fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result { - writeln!(formatter, "{}{}", indent, self.location)?; + match self { + &KdTree::Node { location, ref left, ref right } => { + writeln!(formatter, "{}{}", indent, location)?; + left.print(formatter, &format!("{} ", indent))?; + right.print(formatter, &format!("{} ", indent))?; + }, - if let Some(left) = self.left.as_ref() { - 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(())