From da46959258a8b50a82e7bf1fb3bfdd2a1d89245b Mon Sep 17 00:00:00 2001
From: Thomas Forgione <thomas@forgione.fr>
Date: Tue, 16 Aug 2022 15:40:46 +0200
Subject: [PATCH] Create kd tree by depth

---
 src/example.rs |   2 +-
 src/lib.rs     | 105 ++++++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 91 insertions(+), 16 deletions(-)

diff --git a/src/example.rs b/src/example.rs
index d842146..3b93f1d 100644
--- a/src/example.rs
+++ b/src/example.rs
@@ -31,7 +31,7 @@ fn main() {
         Vector2::new(7.0, 2.0),
     ];
 
-    let kdtree = KdTree::new(
+    let kdtree = KdTree::min_leaf_size(
         &mut points,
         2,
         |coord, ref a, ref b| a.data[coord].partial_cmp(&b.data[coord]).unwrap(),
diff --git a/src/lib.rs b/src/lib.rs
index 3617a02..2d3b29f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,5 @@
-use std::fmt;
 use std::cmp::Ordering;
+use std::fmt;
 
 pub enum Side {
     Left,
@@ -16,18 +16,18 @@ pub enum KdTree<T> {
     Leaf(Vec<T>),
 }
 
-impl<T: Clone> KdTree<T> where {
-    pub fn new<F: Fn(usize, &T, &T) -> Ordering>(
+impl<T: Clone> KdTree<T> {
+    pub fn min_leaf_size<F: Fn(usize, &T, &T) -> Ordering>(
         points: &mut [T],
         dimension: usize,
         comparator: F,
         leaf_size: usize,
     ) -> KdTree<T> {
         // Compute bounding box of points
-        KdTree::new_with_depth(points, dimension, &comparator, leaf_size, 0)
+        KdTree::min_leaf_size_with_depth(points, dimension, &comparator, leaf_size, 0)
     }
 
-    fn new_with_depth<F: Fn(usize, &T, &T) -> Ordering>(
+    fn min_leaf_size_with_depth<F: Fn(usize, &T, &T) -> Ordering>(
         points: &mut [T],
         dimension: usize,
         comparator: &F,
@@ -67,14 +67,14 @@ impl<T: Clone> KdTree<T> where {
         // Recursive call
         KdTree::Node {
             location: median,
-            left: Box::new(KdTree::new_with_depth(
+            left: Box::new(KdTree::min_leaf_size_with_depth(
                 &mut left,
                 dimension,
                 comparator,
                 leaf_size,
                 depth + 1,
             )),
-            right: Box::new(KdTree::new_with_depth(
+            right: Box::new(KdTree::min_leaf_size_with_depth(
                 right,
                 dimension,
                 comparator,
@@ -84,17 +84,90 @@ impl<T: Clone> KdTree<T> where {
         }
     }
 
+    pub fn max_depth<F: Fn(usize, &T, &T) -> Ordering>(
+        points: &mut [T],
+        dimension: usize,
+        comparator: F,
+        max_depth: usize,
+    ) -> KdTree<T> {
+        // Compute bounding box of points
+        KdTree::max_depth_with_depth(points, dimension, &comparator, max_depth, 0)
+    }
+
+    fn max_depth_with_depth<F: Fn(usize, &T, &T) -> Ordering>(
+        points: &mut [T],
+        dimension: usize,
+        comparator: &F,
+        max_depth: usize,
+        depth: usize,
+    ) -> KdTree<T> {
+        // Select axis to work on
+        let axis = depth % dimension;
+
+        // Sort the points by axis
+        points.sort_by(|a, b| comparator(axis, a, b));
+
+        // Find the median
+        let len = points.len();
+
+        // If we reached max depth, make a leaf
+        if depth >= max_depth {
+            let mut contained = vec![];
+            for p in points {
+                contained.push(p.clone());
+            }
+            return KdTree::Leaf(contained);
+        }
+
+        // Else split what's remaining
+        let delimiter = points.len() / 2;
+        let median = points[delimiter].clone();
+
+        // Build the left and right branches
+        let mut left = vec![];
+        for i in &points[0..delimiter + 1] {
+            left.push(i.clone());
+        }
+
+        let right = &mut points[delimiter + 1..len];
+
+        // Recursive call
+        KdTree::Node {
+            location: median,
+            left: Box::new(KdTree::max_depth_with_depth(
+                &mut left,
+                dimension,
+                comparator,
+                max_depth,
+                depth + 1,
+            )),
+            right: Box::new(KdTree::max_depth_with_depth(
+                right,
+                dimension,
+                comparator,
+                max_depth,
+                depth + 1,
+            )),
+        }
+    }
+
     // 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)>) {
-
+    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 } => {
-
+            KdTree::Node {
+                ref location,
+                ref left,
+                ref right,
+            } => {
                 locations.push((location.clone(), Side::Left));
                 left.traverse_leaves_aux(callback, locations);
                 locations.pop();
@@ -102,12 +175,10 @@ impl<T: Clone> KdTree<T> where {
                 locations.push((location.clone(), Side::Right));
                 right.traverse_leaves_aux(callback, locations);
                 locations.pop();
-
-            },
+            }
             KdTree::Leaf(ref elements) => {
                 callback(locations, elements);
             }
-
         }
     }
 }
@@ -115,7 +186,11 @@ impl<T: Clone> KdTree<T> where {
 impl<T: fmt::Display> KdTree<T> {
     fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result {
         match self {
-            &KdTree::Node { ref location, ref left, ref right, } => {
+            &KdTree::Node {
+                ref location,
+                ref left,
+                ref right,
+            } => {
                 writeln!(formatter, "{}{}", indent, location)?;
                 left.print(formatter, &format!("{}  ", indent))?;
                 right.print(formatter, &format!("{}  ", indent))?;