mirror of
https://github.com/andreytkachenko/kdtree-rust.git
synced 2024-11-22 09:26:25 +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> {
|
pub struct Kdtree<KdtreePoint> {
|
||||||
nodes: Vec<KdtreeNode<KdtreePoint>>,
|
nodes: Vec<KdtreeNode<KdtreePoint>>,
|
||||||
|
|
||||||
|
node_adding_dimension: usize
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<KdtreePoint: KdtreePointTrait> Kdtree<KdtreePoint> {
|
impl<KdtreePoint: KdtreePointTrait> Kdtree<KdtreePoint> {
|
||||||
@ -27,6 +29,7 @@ impl<KdtreePoint: KdtreePointTrait> Kdtree<KdtreePoint> {
|
|||||||
|
|
||||||
let mut tree = Kdtree {
|
let mut tree = Kdtree {
|
||||||
nodes: vec![],
|
nodes: vec![],
|
||||||
|
node_adding_dimension: 0,
|
||||||
};
|
};
|
||||||
|
|
||||||
tree.build_tree(&mut points, &rect);
|
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
|
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 {
|
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 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 {
|
} else {
|
||||||
|
current_node.left_node = Some(index_of_new_node);
|
||||||
}
|
|
||||||
|
|
||||||
break;
|
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) {
|
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);
|
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.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> {
|
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
|
closed_found_point
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
fn generate_points(point_count : usize) -> Vec<Point3WithId> {
|
||||||
fn test_against_1000_random_points() {
|
|
||||||
let mut points : Vec<Point3WithId> = vec![];
|
let mut points : Vec<Point3WithId> = vec![];
|
||||||
|
|
||||||
let point_count = 1000usize;
|
|
||||||
for i in 0 .. point_count {
|
for i in 0 .. point_count {
|
||||||
points.push(Point3WithId::new(i as i32, gen_random(),gen_random(),gen_random()));
|
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());
|
let tree = kdtree::kdtree::Kdtree::new(&mut points.clone());
|
||||||
|
|
||||||
//test points pushed into the tree, id should be equal.
|
//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);
|
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