spatialrust_search/
brute.rs

1use crate::{NearestNeighborIndex, Neighbor, RadiusSearchIndex, SpatialIndex};
2
3/// Reference index using brute-force search for correctness tests.
4#[derive(Clone, Debug)]
5pub struct BruteForceIndex {
6    x: Vec<f32>,
7    y: Vec<f32>,
8    z: Vec<f32>,
9}
10
11impl BruteForceIndex {
12    /// Builds a brute-force index from coordinate slices.
13    #[must_use]
14    pub fn from_slices(x: &[f32], y: &[f32], z: &[f32]) -> Self {
15        assert_eq!(x.len(), y.len());
16        assert_eq!(x.len(), z.len());
17        Self { x: x.to_vec(), y: y.to_vec(), z: z.to_vec() }
18    }
19}
20
21impl SpatialIndex for BruteForceIndex {
22    fn len(&self) -> usize {
23        self.x.len()
24    }
25}
26
27impl NearestNeighborIndex for BruteForceIndex {
28    fn nearest_one(&self, x: f32, y: f32, z: f32) -> Option<Neighbor> {
29        brute_force_knn(&self.x, &self.y, &self.z, x, y, z, 1).into_iter().next()
30    }
31
32    fn nearest_k(&self, x: f32, y: f32, z: f32, k: usize) -> Vec<Neighbor> {
33        brute_force_knn(&self.x, &self.y, &self.z, x, y, z, k)
34    }
35}
36
37impl RadiusSearchIndex for BruteForceIndex {
38    fn radius_search(&self, x: f32, y: f32, z: f32, radius: f32) -> Vec<Neighbor> {
39        brute_force_radius(&self.x, &self.y, &self.z, x, y, z, radius)
40    }
41}
42
43/// Finds up to `k` nearest neighbors by brute force.
44#[must_use]
45pub fn brute_force_knn(
46    x: &[f32],
47    y: &[f32],
48    z: &[f32],
49    qx: f32,
50    qy: f32,
51    qz: f32,
52    k: usize,
53) -> Vec<Neighbor> {
54    if k == 0 || x.is_empty() {
55        return Vec::new();
56    }
57
58    let mut neighbors: Vec<Neighbor> = x
59        .iter()
60        .enumerate()
61        .map(|(index, &px)| Neighbor {
62            index,
63            distance_squared: squared_distance(px, y[index], z[index], qx, qy, qz),
64        })
65        .collect();
66    neighbors.sort_by(|a, b| {
67        a.distance_squared.partial_cmp(&b.distance_squared).unwrap_or(std::cmp::Ordering::Equal)
68    });
69    neighbors.truncate(k);
70    neighbors
71}
72
73/// Finds all neighbors within `radius` by brute force.
74#[must_use]
75pub fn brute_force_radius(
76    x: &[f32],
77    y: &[f32],
78    z: &[f32],
79    qx: f32,
80    qy: f32,
81    qz: f32,
82    radius: f32,
83) -> Vec<Neighbor> {
84    let radius_sq = radius * radius;
85    let mut neighbors = Vec::new();
86    for (index, &px) in x.iter().enumerate() {
87        let dist_sq = squared_distance(px, y[index], z[index], qx, qy, qz);
88        if dist_sq <= radius_sq {
89            neighbors.push(Neighbor { index, distance_squared: dist_sq });
90        }
91    }
92    neighbors.sort_by(|a, b| {
93        a.distance_squared.partial_cmp(&b.distance_squared).unwrap_or(std::cmp::Ordering::Equal)
94    });
95    neighbors
96}
97
98fn squared_distance(px: f32, py: f32, pz: f32, qx: f32, qy: f32, qz: f32) -> f32 {
99    let dx = px - qx;
100    let dy = py - qy;
101    let dz = pz - qz;
102    dx * dx + dy * dy + dz * dz
103}
spatialrust_search/brute.rs

spatialrust_search/
brute.rs
