diff --git a/benches/bench.rs b/benches/bench.rs index 0f9f40b..aecdaf3 100644 --- a/benches/bench.rs +++ b/benches/bench.rs @@ -40,3 +40,99 @@ fn bench_nearest_from_kdtree_with_1k_3d_points(b: &mut Bencher) { } b.iter(|| kdtree.nearest(&point.0, 8, &squared_euclidean).unwrap()); } + +#[bench] +fn bench_within_2k_data_01_radius(b: &mut Bencher) { + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within(&point.0, 0.1, &squared_euclidean).unwrap()); +} + +#[bench] +fn bench_within_2k_data_02_radius(b: &mut Bencher) { + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within(&point.0, 0.2, &squared_euclidean).unwrap()); +} + +#[bench] +fn bench_within_unsorted_2k_data_01_radius(b: &mut Bencher) { + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within_unsorted(&point.0, 0.1, &squared_euclidean).unwrap()); +} + +#[bench] +fn bench_within_unsorted_2k_data_02_radius(b: &mut Bencher) { + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within_unsorted(&point.0, 0.2, &squared_euclidean).unwrap()); +} + +#[bench] +fn bench_within_count_2k_data_01_radius(b: &mut Bencher) { + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within_count(&point.0, 0.1, &squared_euclidean).unwrap()); +} + +#[bench] +fn bench_within_count_2k_data_02_radius(b: &mut Bencher) { + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within_count(&point.0, 0.2, &squared_euclidean).unwrap()); +} diff --git a/src/kdtree.rs b/src/kdtree.rs index b2823a3..1410b0c 100644 --- a/src/kdtree.rs +++ b/src/kdtree.rs @@ -93,14 +93,11 @@ impl + std::cmp::Pa .collect()) } - pub fn within(&self, point: &[A], radius: A, distance: &F) -> Result, ErrorKind> + #[inline(always)] + fn evaluated_heap(&self, point: &[A], radius: A, distance: &F) -> BinaryHeap> where F: Fn(&[A], &[A]) -> A, { - self.check_point(point)?; - if self.size == 0 { - return Ok(vec![]); - } let mut pending = BinaryHeap::new(); let mut evaluated = BinaryHeap::>::new(); pending.push(HeapElement { @@ -110,9 +107,45 @@ impl + std::cmp::Pa while !pending.is_empty() && (-pending.peek().unwrap().distance <= radius) { self.nearest_step(point, self.size, radius, distance, &mut pending, &mut evaluated); } + evaluated + } + + pub fn within(&self, point: &[A], radius: A, distance: &F) -> Result, ErrorKind> + where + F: Fn(&[A], &[A]) -> A, + { + self.check_point(point)?; + if self.size == 0 { + return Ok(vec![]); + } + let evaluated = self.evaluated_heap(point, radius, distance); Ok(evaluated.into_sorted_vec().into_iter().map(Into::into).collect()) } + pub fn within_unsorted(&self, point: &[A], radius: A, distance: &F) -> Result, ErrorKind> + where + F: Fn(&[A], &[A]) -> A, + { + self.check_point(point)?; + if self.size == 0 { + return Ok(vec![]); + } + let evaluated = self.evaluated_heap(point, radius, distance); + Ok(evaluated.into_iter().map(Into::into).collect()) + } + + pub fn within_count(&self, point: &[A], radius: A, distance: &F) -> Result + where + F: Fn(&[A], &[A]) -> A, + { + self.check_point(point)?; + if self.size == 0 { + return Ok(0); + } + let evaluated = self.evaluated_heap(point, radius, distance); + Ok(evaluated.len()) + } + fn nearest_step<'b, F>( &self, point: &[A], diff --git a/tests/kdtree.rs b/tests/kdtree.rs index 16a6984..ee1f23b 100644 --- a/tests/kdtree.rs +++ b/tests/kdtree.rs @@ -60,6 +60,39 @@ fn it_works() { vec![(0.0, &1), (2.0, &2), (2.0, &0)] ); + let unsorted1 = kdtree.within_unsorted(&POINT_A.0, 0.0, &squared_euclidean).unwrap(); + let ans1 = vec![(0.0, &0)]; + assert_eq!(unsorted1.len(), ans1.len()); + assert_eq!( + kdtree.within_count(&POINT_A.0, 0.0, &squared_euclidean).unwrap(), + ans1.len() + ); + for item in unsorted1 { + assert!(ans1.contains(&item)); + } + + let unsorted2 = kdtree.within_unsorted(&POINT_B.0, 1.0, &squared_euclidean).unwrap(); + let ans2 = vec![(0.0, &1)]; + assert_eq!(unsorted2.len(), ans2.len()); + assert_eq!( + kdtree.within_count(&POINT_B.0, 1.0, &squared_euclidean).unwrap(), + ans2.len() + ); + for item in unsorted2 { + assert!(ans2.contains(&item)); + } + + let unsorted3 = kdtree.within_unsorted(&POINT_B.0, 2.0, &squared_euclidean).unwrap(); + let ans3 = vec![(0.0, &1), (2.0, &2), (2.0, &0)]; + assert_eq!(unsorted3.len(), ans3.len()); + assert_eq!( + kdtree.within_count(&POINT_B.0, 2.0, &squared_euclidean).unwrap(), + ans3.len() + ); + for item in unsorted3 { + assert!(ans3.contains(&item)); + } + assert_eq!( kdtree .iter_nearest(&POINT_A.0, &squared_euclidean)