Cleaner, min leaf size

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);
 

src/lib.rs

@@ -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(())