space/
linear.rs

1use core::{iter::Map, slice};
2
3use alloc::vec::{self, Vec};
4use pgat::{ProxyView, ReferenceProxy, View};
5
6use crate::{ApproximateSpace, ExactSpace, Knn, Metric, SpatialContainer, linear_nn};
7
8/// This function performs exact linear nearest neighbor search.
9///
10/// This may be useful specifically when implementing spatial containers
11/// where you need to abstract over ProxyView types.
12pub fn linear_knn<'a, 'b, M, P, V>(
13    metric: M,
14    dataset: impl Iterator<Item = (View<'a, P>, View<'a, V>)>,
15    query: View<'b, P>,
16    num: usize,
17) -> Vec<(M::Unit, View<'a, P>, View<'a, V>)>
18where
19    M: Metric<P>,
20    P: ProxyView,
21    V: ProxyView,
22{
23    let mut dataset = dataset.map(|(pt, val)| (metric.distance(pt, query), pt, val));
24
25    // Create a vector with the correct capacity in advance.
26    let mut neighbors = Vec::with_capacity(num);
27
28    // Extend the vector with the first `num` neighbors.
29    neighbors.extend((&mut dataset).take(num));
30    // Sort the vector by the neighbor distance.
31    neighbors.sort_unstable_by_key(|n| n.0);
32
33    // Iterate over each additional neighbor.
34    for point in dataset {
35        // Find the position at which it would be inserted.
36        let position = neighbors.partition_point(|n| n.0 <= point.0);
37        // If the point is closer than at least one of the points already in `neighbors`, add it
38        // into its sorted position.
39        if position != num {
40            neighbors.pop();
41            neighbors.insert(position, point);
42        }
43    }
44
45    neighbors
46}
47
48/// Performs a linear knn search by iterating one-by-one over the dataset
49/// and keeping a running set of neighbors which it searches through with binary search.
50///
51/// You may use the optional type parameters `PP` and `VP` to specify the proxy types for the point and value.
52/// By default, it uses [`ReferenceProxy`] for both point and value, which uses &P and &V as the proxies.
53pub struct LinearSearch<'a, M, P, V, PP = ReferenceProxy<P>, VP = ReferenceProxy<V>> {
54    pub metric: M,
55    pub data: &'a [(P, V)],
56    pub _phantom: core::marker::PhantomData<(PP, VP)>,
57}
58
59impl<'a, M, P, V, PP, VP> LinearSearch<'a, M, P, V, PP, VP>
60where
61    M: Metric<PP>,
62    PP: ProxyView<Owned = P>,
63    VP: ProxyView<Owned = V>,
64{
65    pub fn new(metric: M, data: &'a [(P, V)]) -> Self {
66        Self {
67            metric,
68            data,
69            _phantom: core::marker::PhantomData,
70        }
71    }
72}
73
74impl<'a, M, P, V, PP, VP> ApproximateSpace for LinearSearch<'a, M, P, V, PP, VP>
75where
76    M: Metric<PP>,
77    PP: ProxyView<Owned = P>,
78    VP: ProxyView<Owned = V>,
79{
80    type PointProxy = PP;
81    type ValueProxy = VP;
82    type Metric = M;
83}
84
85/// This trait is implemented for linear search, which is an exact search algorithm.
86impl<'a, M, P, V, PP, VP> ExactSpace for LinearSearch<'a, M, P, V, PP, VP>
87where
88    M: Metric<PP>,
89    PP: ProxyView<Owned = P>,
90    VP: ProxyView<Owned = V>,
91{
92}
93
94impl<'c, M, P, V, PP, VP> Knn for LinearSearch<'c, M, P, V, PP, VP>
95where
96    M: Metric<PP>,
97    PP: ProxyView<Owned = P>,
98    VP: ProxyView<Owned = V>,
99{
100    type KnnIter<'a>
101        = vec::IntoIter<(M::Unit, View<'a, PP>, View<'a, VP>)>
102    where
103        Self: 'a;
104
105    fn knn<'a, 'b>(&'a self, query: View<'b, Self::PointProxy>, num: usize) -> Self::KnnIter<'a> {
106        linear_knn::<M, PP, VP>(
107            self.metric,
108            self.data
109                .iter()
110                .map(|(pt, val)| (PP::view(pt), VP::view(val))),
111            query,
112            num,
113        )
114        .into_iter()
115    }
116
117    fn nn<'a, 'b>(
118        &'a self,
119        query: View<'b, Self::PointProxy>,
120    ) -> Option<(M::Unit, View<'a, PP>, View<'a, VP>)> {
121        linear_nn::<M, PP, VP>(
122            self.metric,
123            self.data
124                .iter()
125                .map(|(pt, val)| (PP::view(pt), VP::view(val))),
126            query,
127        )
128    }
129}
130
131pub struct LinearContainer<M, P, V, PP = ReferenceProxy<P>, VP = ReferenceProxy<V>> {
132    pub metric: M,
133    pub data: Vec<(P, V)>,
134    pub _phantom: core::marker::PhantomData<(PP, VP)>,
135}
136
137impl<M, P, V, PP, VP> ApproximateSpace for LinearContainer<M, P, V, PP, VP>
138where
139    M: Metric<PP>,
140    PP: ProxyView<Owned = P>,
141    VP: ProxyView<Owned = V>,
142{
143    type PointProxy = PP;
144    type ValueProxy = VP;
145    type Metric = M;
146}
147
148/// This trait is implemented for linear search, which is an exact search algorithm.
149impl<M, P, V, PP, VP> ExactSpace for LinearContainer<M, P, V, PP, VP>
150where
151    M: Metric<PP>,
152    PP: ProxyView<Owned = P>,
153    VP: ProxyView<Owned = V>,
154{
155}
156
157impl<M, P, V, PP, VP> SpatialContainer for LinearContainer<M, P, V, PP, VP>
158where
159    M: Metric<PP>,
160    PP: ProxyView<Owned = P>,
161    VP: ProxyView<Owned = V>,
162{
163    type SpatialIter<'a>
164        = Map<slice::Iter<'a, (P, V)>, fn(&'a (P, V)) -> (View<'a, PP>, View<'a, VP>)>
165    where
166        Self: 'a;
167
168    fn with_metric(metric: Self::Metric) -> Self {
169        Self {
170            metric,
171            data: Vec::new(),
172            _phantom: core::marker::PhantomData,
173        }
174    }
175
176    fn insert(&mut self, point: P, value: V) {
177        self.data.push((point, value));
178    }
179
180    fn iter(&self) -> Self::SpatialIter<'_> {
181        self.data.iter().map(|(p, v)| (PP::view(p), VP::view(v)))
182    }
183}
184
185impl<M, P, V, PP, VP> Knn for LinearContainer<M, P, V, PP, VP>
186where
187    M: Metric<PP>,
188    PP: ProxyView<Owned = P>,
189    VP: ProxyView<Owned = V>,
190{
191    type KnnIter<'a>
192        = vec::IntoIter<(M::Unit, View<'a, PP>, View<'a, VP>)>
193    where
194        Self: 'a;
195
196    fn knn<'a, 'b>(&'a self, query: View<'b, Self::PointProxy>, num: usize) -> Self::KnnIter<'a> {
197        linear_knn::<M, PP, VP>(
198            self.metric,
199            self.data
200                .iter()
201                .map(|(pt, val)| (PP::view(pt), VP::view(val))),
202            query,
203            num,
204        )
205        .into_iter()
206    }
207
208    fn nn<'a, 'b>(
209        &'a self,
210        query: View<'b, Self::PointProxy>,
211    ) -> Option<(M::Unit, View<'a, PP>, View<'a, VP>)> {
212        linear_nn::<M, PP, VP>(
213            self.metric,
214            self.data
215                .iter()
216                .map(|(pt, val)| (PP::view(pt), VP::view(val))),
217            query,
218        )
219    }
220}
space/linear.rs

space/
linear.rs
