Implemented insert node, need to make it so that tree sometimes automatically rebalances itself

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,18 +52,33 @@ 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;
+                }
+            }
+        }
     }
 
     fn nearest_search_impl(&self, p: &KdtreePoint, searched_index: usize, best_distance_squared: &mut f64, best_leaf_found: &mut usize) {
@@ -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> {

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.
@@ -81,3 +87,33 @@ 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);
+    }
+}