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

This commit is contained in:
Olek 2016-12-30 16:54:50 +01:00
parent b236cd2595
commit 5b195d3891
2 changed files with 72 additions and 11 deletions

View File

@ -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> {

View File

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