added within_unsorted and within_count (#52)

* added within_unsorted within_count * did cargo fmt * better org --------- Co-authored-by: Rui Hu <[email protected]>
mrhooray · Feb 2, 2024 · 4f3755e · 4f3755e
1 parent 5e01eb8
commit 4f3755e
Show file tree

Hide file tree

Showing 3 changed files with 167 additions and 5 deletions.
diff --git 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
@@ -93,14 +93,11 @@ impl<A: Float + Zero + One, T: std::cmp::PartialEq, U: AsRef<[A]> + std::cmp::Pa
             .collect())
     }
 
-    pub fn within<F>(&self, point: &[A], radius: A, distance: &F) -> Result<Vec<(A, &T)>, ErrorKind>
+    #[inline(always)]
+    fn evaluated_heap<F>(&self, point: &[A], radius: A, distance: &F) -> BinaryHeap<HeapElement<A, &T>>
     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::<HeapElement<A, &T>>::new();
         pending.push(HeapElement {
@@ -110,9 +107,45 @@ impl<A: Float + Zero + One, T: std::cmp::PartialEq, U: AsRef<[A]> + 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<F>(&self, point: &[A], radius: A, distance: &F) -> Result<Vec<(A, &T)>, 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<F>(&self, point: &[A], radius: A, distance: &F) -> Result<Vec<(A, &T)>, 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<F>(&self, point: &[A], radius: A, distance: &F) -> Result<usize, ErrorKind>
+    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
@@ -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)