add distance to results

This commit is contained in:
Andrey Tkachenko 2018-12-10 18:54:14 +04:00
parent caa2e5cbf8
commit b67df3d4a1
2 changed files with 19 additions and 18 deletions

View File

@ -32,7 +32,7 @@ pub struct NearestNeighboursIter<'a, 'b, T> {
impl<'a, 'b, T> Iterator for NearestNeighboursIter<'a, 'b, T>
where T: KdTreePoint
{
type Item = &'a T;
type Item = (f64, &'a T);
fn next(&mut self) -> Option<Self::Item> {
let p = self.ref_node;
@ -54,8 +54,9 @@ impl<'a, 'b, T> Iterator for NearestNeighboursIter<'a, 'b, T>
self.node_stack.push(idx);
}
if p.dist(&node.point) <= self.range {
return Some(&node.point);
let dist = p.dist(&node.point);
if dist <= self.range {
return Some((dist, &node.point));
}
} else if point_splitting_dim_value <= splitting_value {
if let Some(idx) = node.left_node {
@ -115,13 +116,13 @@ impl<KP: KdTreePoint> KdTree<KP> {
self.rebuild_tree(&mut points);
}
pub fn nearest_search(&self, node: &KP) -> KP {
pub fn nearest_search(&self, node: &KP) -> (f64, &KP) {
let mut nearest_neighbor = 0usize;
let mut best_distance = self.nodes[0].point.dist(&node);
self.nearest_search_impl(node, 0usize, &mut best_distance, &mut nearest_neighbor);
self.nodes[nearest_neighbor].point
(best_distance, &self.nodes[nearest_neighbor].point)
}
pub fn nearest_search_dist<'a, 'b>(&'a self, node: &'b KP, dist: f64) -> NearestNeighboursIter<'a, 'b, KP> {
@ -143,7 +144,7 @@ impl<KP: KdTreePoint> KdTree<KP> {
}
pub fn distance_squared_to_nearest(&self, node: &KP) -> f64 {
self.nearest_search(node).dist(&node)
self.nearest_search(node).0
}
pub fn insert_nodes_and_rebuild(&mut self, nodes_to_add : &mut [KP]) {
@ -354,9 +355,9 @@ mod tests {
let tree = KdTree::new(&mut point_vec.clone());
for p in &point_vec {
let found_nn = tree.nearest_search(p);
let found_nn = tree.nearest_search(p).1;
assert_eq!(p.id,found_nn.id);
assert_eq!(p.id, found_nn.id);
}
true

View File

@ -11,7 +11,7 @@ fn gen_random() -> f64 {
rand::thread_rng().gen_range(0., 1000.)
}
fn find_nn_with_linear_search(points : &Vec<Point3WithId>, find_for : Point3WithId) -> &Point3WithId {
fn find_nn_with_linear_search(points : &Vec<Point3WithId>, find_for : Point3WithId) -> (f64, &Point3WithId) {
let mut best_found_distance = squared_euclidean(find_for.dims(), points[0].dims());
let mut closed_found_point = &points[0];
@ -24,17 +24,17 @@ fn find_nn_with_linear_search(points : &Vec<Point3WithId>, find_for : Point3With
}
}
closed_found_point
(best_found_distance, closed_found_point)
}
fn find_neigbours_with_linear_search(points : &Vec<Point3WithId>, find_for : Point3WithId, dist: f64) -> Vec<&Point3WithId> {
fn find_neigbours_with_linear_search(points : &Vec<Point3WithId>, find_for : Point3WithId, dist: f64) -> Vec<(f64, &Point3WithId)> {
let mut result = Vec::new();
for p in points {
let d = squared_euclidean(find_for.dims(), p.dims());
if d <= dist {
result.push(p);
result.push((d, p));
}
}
@ -64,7 +64,7 @@ fn test_against_1000_random_points() {
for i in 0 .. point_count {
let p = &points[i];
assert_eq!(p.id, tree.nearest_search(p).id );
assert_eq!(p.id, tree.nearest_search(p).1.id );
}
//test randomly generated points within the cube. and do the linear search. should match
@ -74,7 +74,7 @@ fn test_against_1000_random_points() {
let found_by_linear_search = find_nn_with_linear_search(&points, p);
let point_found_by_kdtree = tree.nearest_search(&p);
assert_eq!(point_found_by_kdtree.id, found_by_linear_search.id);
assert_eq!(point_found_by_kdtree.1.id, found_by_linear_search.1.id);
}
}
@ -97,14 +97,14 @@ fn test_incrementally_build_tree_against_built_at_once() {
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);
assert_eq!(tree_built_at_once.nearest_search(p).1.id, tree_built_incrementally.nearest_search(p).1.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);
assert_eq!(tree_built_at_once.nearest_search(&p).1.id, tree_built_incrementally.nearest_search(&p).1.id);
}
}
@ -125,8 +125,8 @@ fn test_neighbour_search_with_distance() {
assert_eq!(found_by_linear_search.len(), point_found_by_kdtree.len());
if point_found_by_kdtree.len() > 0 {
found_by_linear_search.sort_by(|a, b| a.id.cmp(&b.id));
point_found_by_kdtree.sort_by(|a, b| a.id.cmp(&b.id));
found_by_linear_search.sort_by(|a, b| a.1.id.cmp(&b.1.id));
point_found_by_kdtree.sort_by(|a, b| a.1.id.cmp(&b.1.id));
}
assert_eq!(point_found_by_kdtree, found_by_linear_search);