mirror of
https://github.com/andreytkachenko/kdtree-rust.git
synced 2024-11-21 17:06:24 +04:00
Implemented insert node, need to make it so that tree sometimes automatically rebalances itself
This commit is contained in:
parent
b236cd2595
commit
5b195d3891
@ -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> {
|
||||
|
@ -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);
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue
Block a user