From 54655dd97114edb269466570879e388be1811f69 Mon Sep 17 00:00:00 2001
From: Joe <Joe>
Date: Sat, 18 Oct 2025 13:03:11 -0700
Subject: [PATCH] =?UTF-8?q?add=20nearest=20search,=20takes=2070=C2=B5s?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 benches/benchmarks.rs |  19 +++++++-
 src/lib.rs            | 105 +++++++++++++++++++++++++++++++++++++++---
 2 files changed, 117 insertions(+), 7 deletions(-)

diff --git a/benches/benchmarks.rs b/benches/benchmarks.rs
index 6ff7834..e9244ac 100644
--- a/benches/benchmarks.rs
+++ b/benches/benchmarks.rs
@@ -60,5 +60,22 @@ fn update_positions(c: &mut Criterion) {
     });
 }
 
-criterion_group!(benches, new_index, update_positions);
+fn nearest(c: &mut Criterion) {
+    c.bench_function("nearst point", move |b| {
+        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1)
+            .into_iter()
+            .map(|[x, y]| Point2d { x, y })
+            .collect();
+
+        let index = Spindex2d::new(&points);
+
+        let mut i = 0;
+        b.iter(|| {
+            index.nearest(&points[i]);
+            i += 1;
+        });
+    });
+}
+
+criterion_group!(benches, new_index, update_positions, nearest);
 criterion_main!(benches);
diff --git a/src/lib.rs b/src/lib.rs
index 8350f3e..df8fd94 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,4 @@
-use std::{sync::RwLock, thread};
+use std::{collections::HashMap, sync::RwLock, thread};
 
 #[derive(Debug, Default, Clone, Copy, PartialEq, PartialOrd)]
 pub struct Point2d {
@@ -6,6 +6,20 @@ pub struct Point2d {
     pub y: f64,
 }
 
+impl Point2d {
+    pub fn new(x: f64, y: f64) -> Self {
+        Self { x, y }
+    }
+
+    pub fn distance_squared(&self, other: &Self) -> f64 {
+        (self.x - other.x).powi(2) + (self.y - other.y).powi(2)
+    }
+
+    pub fn distance(&self, other: &Self) -> f64 {
+        self.distance_squared(other).sqrt()
+    }
+}
+
 #[derive(Debug, Default, Clone, Copy)]
 pub struct PVal {
     point_index: usize,
@@ -28,9 +42,7 @@ impl PartialOrd for PVal {
 
 impl Ord for PVal {
     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-        self.val
-            .partial_cmp(&other.val)
-            .unwrap_or(std::cmp::Ordering::Equal)
+        self.val.total_cmp(&other.val)
     }
 }
 
@@ -55,12 +67,14 @@ impl PVal {
 pub struct Spindex2d {
     pub xs: RwLock<Vec<PVal>>,
     pub ys: RwLock<Vec<PVal>>,
+    len: usize,
 }
 
 impl Spindex2d {
     pub fn new(points: &[Point2d]) -> Self {
-        let mut xs = Vec::with_capacity(points.len());
-        let mut ys = Vec::with_capacity(points.len());
+        let len = points.len();
+        let mut xs = Vec::with_capacity(len);
+        let mut ys = Vec::with_capacity(len);
 
         thread::scope(|s| {
             s.spawn(|| {
@@ -81,6 +95,7 @@ impl Spindex2d {
         let index = Self {
             xs: xs.into(),
             ys: ys.into(),
+            len,
         };
 
         index.sort();
@@ -111,6 +126,84 @@ impl Spindex2d {
         self.sort();
     }
 
+    /// Find up to `k` closest points to `point`.
+    pub fn knn(&self, point: &Point2d, k: usize) -> Option<Vec<Point2d>> {
+        todo!()
+    }
+
+    /// Search within a given `radius` from `point`.
+    pub fn radius(&self, point: &Point2d, radius: f64) -> Option<Vec<Point2d>> {
+        todo!()
+    }
+
+    /// Search an axis-aligned bounding box whose sides are `span` away from the given point.
+    pub fn aabb(&self, point: &Point2d, span: f64) -> Option<Vec<Point2d>> {
+        todo!()
+    }
+
+    pub fn nearest(&self, point: &Point2d) -> Option<Point2d> {
+        let x = point.x;
+        let y = point.y;
+
+        let xs = self.xs.read().unwrap();
+        let ys = self.ys.read().unwrap();
+
+        let mut nxs = HashMap::with_capacity(100);
+        let mut nys = HashMap::with_capacity(100);
+        thread::scope(|s| {
+            s.spawn(|| {
+                let p = PVal::new(0, x);
+                let nx = match xs.binary_search(&p) {
+                    Ok(i) | Err(i) => i,
+                };
+                let start = nx.saturating_sub(50);
+                let end = (nx + (50)).min(self.len);
+                for p in xs[start..end].iter() {
+                    nxs.insert(p.point_index(), p.val());
+                }
+            });
+
+            s.spawn(|| {
+                let p = PVal::new(0, y);
+                let ny = match ys.binary_search(&p) {
+                    Ok(i) | Err(i) => i,
+                };
+                let start = ny.saturating_sub(50);
+                let end = (ny + 50).min(self.len);
+                for p in ys[start..end].iter() {
+                    nys.insert(p.point_index(), p.val());
+                }
+            });
+        });
+        let mut points = Vec::with_capacity(100);
+        for (id, xval) in nxs.into_iter() {
+            if let Some(&yval) = nys.get(&id) {
+                points.push(Point2d::new(xval, yval));
+            }
+        }
+        if points.is_empty() {
+            None
+        } else {
+            points.sort_unstable_by(|a, b| {
+                let adist = a.distance_squared(point);
+                let bdist = b.distance_squared(point);
+                adist.total_cmp(&bdist)
+            });
+            points.first().cloned()
+        }
+    }
+
+    pub fn points(&self) -> Vec<Point2d> {
+        let mut points = vec![Point2d::default(); self.xs.read().unwrap().len()];
+        for x in self.xs.read().unwrap().iter() {
+            points[x.point_index].x = x.val;
+        }
+        for y in self.ys.read().unwrap().iter() {
+            points[y.point_index].y = y.val;
+        }
+        points
+    }
+
     fn sort(&self) {
         thread::scope(|s| {
             s.spawn(|| {