From 5b195d3891f44920b2b0d865bedede105e098aee Mon Sep 17 00:00:00 2001
From: Olek <ntszar@gmail.com>
Date: Fri, 30 Dec 2016 16:54:50 +0100
Subject: [PATCH] Implemented insert node, need to make it so that tree
 sometimes automatically rebalances itself

---
 src/kdtree/mod.rs          | 41 +++++++++++++++++++++++++++++--------
 tests/integration_tests.rs | 42 +++++++++++++++++++++++++++++++++++---
 2 files changed, 72 insertions(+), 11 deletions(-)

diff --git a/src/kdtree/mod.rs b/src/kdtree/mod.rs
index 050eba8..5391f41 100644
--- a/src/kdtree/mod.rs
+++ b/src/kdtree/mod.rs
@@ -15,6 +15,8 @@ pub trait KdtreePointTrait: Copy {
 
 pub struct Kdtree<KdtreePoint> {
     nodes: Vec<KdtreeNode<KdtreePoint>>,
+
+    node_adding_dimension: usize
 }
 
 impl<KdtreePoint: KdtreePointTrait> Kdtree<KdtreePoint> {
@@ -27,6 +29,7 @@ impl<KdtreePoint: KdtreePointTrait> Kdtree<KdtreePoint> {
 
         let mut tree = Kdtree {
             nodes: vec![],
+            node_adding_dimension: 0,
         };
 
         tree.build_tree(&mut points, &rect);
@@ -49,17 +52,32 @@ impl<KdtreePoint: KdtreePointTrait> Kdtree<KdtreePoint> {
         squared_euclidean(&self.nearest_search(node).dims(), node.dims()) <= squared_range
     }
 
-    pub fn insert_node(&mut self, node_to_add : &KdtreePoint) {
+    pub fn insert_node(&mut self, node_to_add : KdtreePoint) {
+
+        let mut current_index = 0;
+        let dimension = self.node_adding_dimension;
+        let index_of_new_node = self.add_node(node_to_add,dimension,node_to_add.dims()[dimension]);
+        self.node_adding_dimension = ( dimension + 1) % node_to_add.dims().len();
 
-        let current_index = 0;
         loop {
-            let current_node = &self.nodes[current_index];
+
+            let current_node = &mut self.nodes[current_index];
+
             if node_to_add.dims()[current_node.dimension] <= current_node.split_on {
+                if let Some(left_node_index) = current_node.left_node {
+                    current_index = left_node_index
+                } else {
+                    current_node.left_node = Some(index_of_new_node);
+                    break;
+                }
             } else {
-
+                if let Some(right_node_index) = current_node.right_node {
+                    current_index = right_node_index
+                } else {
+                    current_node.right_node = Some(index_of_new_node);
+                    break;
+                }
             }
-
-            break;
         }
     }
 
@@ -234,10 +252,17 @@ mod tests {
 
         let mut tree = Kdtree::new(&mut vec);
 
-        tree.insert_node(&Point2WithId::new(0,1.,0.));
-        tree.insert_node(&Point2WithId::new(0,-1.,0.));
+        tree.insert_node(Point2WithId::new(0,1.,0.));
+        tree.insert_node(Point2WithId::new(0,-1.,0.));
 
         assert_eq!(tree.nodes.len(), 3);
+        assert_eq!(tree.nodes[0].dimension, 0);
+
+        assert_eq!(tree.nodes[0].left_node.is_some(), true);
+        assert_eq!(tree.nodes[1].point.dims()[0], 1.);
+        assert_eq!(tree.nodes[2].point.dims()[0], -1.);
+
+        assert_eq!(tree.nodes[0].right_node.is_some(), true);
     }
 
     fn qc_value_vec_to_2d_points_vec(xs: &Vec<f64>) -> Vec<Point2WithId> {
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index b8735a3..b1f9051 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -53,15 +53,21 @@ fn find_nn_with_linear_search<'a>(points : &'a Vec<Point3WithId>, find_for : Poi
     closed_found_point
 }
 
-#[test]
-fn test_against_1000_random_points() {
+fn generate_points(point_count : usize) -> Vec<Point3WithId> {
     let mut points : Vec<Point3WithId> = vec![];
 
-    let point_count = 1000usize;
     for i in 0 .. point_count {
         points.push(Point3WithId::new(i as i32, gen_random(),gen_random(),gen_random()));
     }
 
+    points
+}
+
+#[test]
+fn test_against_1000_random_points() {
+    let point_count = 1000usize;
+    let mut points = generate_points(point_count);
+
     let tree = kdtree::kdtree::Kdtree::new(&mut points.clone());
 
     //test points pushed into the tree, id should be equal.
@@ -80,4 +86,34 @@ fn test_against_1000_random_points() {
 
         assert_eq!(point_found_by_kdtree.id, found_by_linear_search.id);
     }
+}
+
+#[test]
+fn test_incrementally_build_tree_against_built_at_once() {
+    let point_count = 2000usize;
+    let mut points = generate_points(point_count);
+
+    let tree_built_at_once = kdtree::kdtree::Kdtree::new(&mut points.clone());
+    let mut tree_built_incrementally = kdtree::kdtree::Kdtree::new(&mut points[0..1]);
+
+    for i in 1 .. point_count {
+        let p = &points[i];
+
+        tree_built_incrementally.insert_node(p.clone());
+    }
+
+
+    //test points pushed into the tree, id should be equal.
+    for i in 0 .. point_count {
+        let p = &points[i];
+
+        assert_eq!(tree_built_at_once.nearest_search(p).id, tree_built_incrementally.nearest_search(p).id);
+    }
+
+
+    //test randomly generated points within the cube. and do the linear search. should match
+    for _ in 0 .. 5000 {
+        let p = Point3WithId::new(0i32, gen_random(), gen_random(), gen_random());
+        assert_eq!(tree_built_at_once.nearest_search(&p).id, tree_built_incrementally.nearest_search(&p).id);
+    }
 }
\ No newline at end of file