1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
use std::{cmp::Ordering, collections::BinaryHeap};
use rand::Rng;
use num_traits::Float;
/// A node of the vantage point tree.
#[derive(Clone, Debug)]
pub(crate) struct Node<T: Float> {
index: usize,
threshold: T,
left: Option<Box<Node<T>>>,
right: Option<Box<Node<T>>>,
}
impl<T: Float> Default for Node<T> {
fn default() -> Self {
Node {
index: 0,
threshold: T::zero(),
left: None,
right: None,
}
}
}
/// An item on the intermediate result heap. It is used to store results from the nearest neighbors
/// search performed of the vantage point tree.
struct HeapItem<T: Float> {
// Index of a sample.
index: usize,
// Distance of the sample from the target.
distance: T,
}
impl<T: Float> PartialOrd for HeapItem<T> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl<T: Float> PartialEq for HeapItem<T> {
fn eq(&self, other: &Self) -> bool {
self.distance == other.distance
}
}
impl<T: Float> Eq for HeapItem<T> {}
impl<T: Float> Ord for HeapItem<T> {
fn cmp(&self, other: &Self) -> Ordering {
self.distance.partial_cmp(&other.distance).unwrap()
}
}
/// Auxiliary struct used to build the vantage point tree.
struct VPTreeBuilder<'a, T: Float> {
root: &'a mut Option<Box<Node<T>>>,
lower: usize,
upper: usize,
}
impl<'a, T: Float> VPTreeBuilder<'a, T> {
/// VPTreeBuilder constructor.
///
/// # Arguments
///
/// * `node` - current node.
///
/// * `lower` - lower bound.
///
/// * `upper` - upper bound.
fn new(root: &'a mut Option<Box<Node<T>>>, lower: usize, upper: usize) -> Self {
Self { root, lower, upper }
}
}
/// Vantage Point tree.
pub(crate) struct VPTree<'a, T: Float + Send + Sync, U> {
items: Vec<(usize, &'a U)>,
pub(crate) root: Option<Box<Node<T>>>,
}
impl<'a, T: Float + Send + Sync, U> VPTree<'a, T, U> {
/// Constructor for the `VPTree` struct.
///
/// # Arguments
///
/// * `items` - **original** items to build the tree on.
///
/// * `metric_f` - metric function.
pub fn new<F>(items: &'a [U], metric_f: F) -> Self
where
F: Fn(&U, &U) -> T,
{
let mut tree = VPTree {
// Need to swap some references around, don't want to move original data.
// Also need to keep track of the original position, as it will differ.
items: items.iter().enumerate().collect(),
root: None,
};
let n_samples = tree.items.len(); // Immutable borrow must be kept outside.
tree.build_from_points(0, n_samples, &metric_f);
tree
}
/// Function that iteratively fills the tree.
///
/// # Arguments
///
/// * `lower' - lower bound.
///
/// * `upper` - upper bound.
///
/// * `metric_f` - metric function.
fn build_from_points<F>(&mut self, lower: usize, upper: usize, metric_f: F)
where
F: Fn(&U, &U) -> T,
{
let mut stack = vec![VPTreeBuilder::new(&mut self.root, lower, upper)];
let mut thread_rng = super::make_rng();
while let Some(builder) = stack.pop() {
let VPTreeBuilder { root, lower, upper } = builder;
// Lower index is center of current node.
if upper != lower {
*root = Some(Box::new(Node::default()));
let node = root.as_deref_mut().unwrap();
node.index = lower;
if upper - lower > 1 {
// If we did not arrive at leaf yet choose an arbitrary point
// and move it to the start.
let i = thread_rng.random_range(lower..upper);
self.items.swap(lower, i);
let (_, to_cmp) = self.items[lower];
// Partition around the median distances.
let median: usize = (upper + lower) / 2;
self.items[lower + 1..upper].select_nth_unstable_by(
median - lower - 1,
&mut |(_, a): &(usize, &U), (_, b): &(usize, &U)| {
if metric_f(to_cmp, a) < metric_f(to_cmp, b) {
Ordering::Less
} else if metric_f(to_cmp, a) == metric_f(to_cmp, b) {
Ordering::Equal
} else {
Ordering::Greater
}
},
);
// Threshold of the new node will be the distances to the median.
node.threshold = metric_f(self.items[lower].1, self.items[median].1);
stack.push(VPTreeBuilder::new(&mut node.left, lower + 1, median));
stack.push(VPTreeBuilder::new(&mut node.right, median, upper));
}
}
}
}
/// Auxiliary function that searches for the k nearest neighbors of an item.
fn look_up<F>(
&self,
tau: &mut T, // Tracks the distances to the farthest point in the results.
target: &U,
k: usize,
heap: &mut BinaryHeap<HeapItem<T>>,
metric_f: F,
) where
F: Fn(&U, &U) -> T,
{
let mut stack: Vec<&Option<Box<Node<T>>>> = vec![&self.root];
while let Some(next_in_stack) = stack.pop() {
if let Some(node) = next_in_stack {
let (original_position, point) = &self.items[node.index];
// Compute distances between target and current node.
let distance: T = metric_f(point, target);
if distance < *tau {
// If current node is within the radius tau
// remove furthest node from result list (if it already contains k results),
// add current node to result list and
// update value of tau (farthest point in result list).
if heap.len() == k {
heap.pop();
}
heap.push(HeapItem {
index: *original_position,
distance,
});
if heap.len() == k {
*tau = heap.peek().unwrap().distance;
}
}
match (node.left.as_ref(), node.right.as_ref()) {
// Return if we arrived at a leaf.
(None, None) => continue,
(_, _) => {
// If the target lies within the radius of ball.
if distance < node.threshold {
// if there can still be neighbors inside the ball, recursively search left child first.
if distance - *tau <= node.threshold {
stack.push(&node.left)
}
// if there can still be neighbors outside the ball, recursively search right child.
if distance + *tau >= node.threshold {
stack.push(&node.right)
}
} else {
// If the target lies outsize the radius of the ball.
// If there can still be neighbors outside the ball, recursively search right child.
if distance + *tau >= node.threshold {
stack.push(&node.right)
}
// if there can still be neighbors inside the ball, recursively search left child.
if distance - *tau <= node.threshold {
stack.push(&node.left)
}
}
}
}
}
}
}
/// Function that searches the tree and finds the k nearest neighbors of `target`.
///
/// # Arguments
///
/// * `target` - target data point.
///
/// * `index` - index of the target.
///
/// * `k` - number of nearest neighbors.
///
/// * `results` - vector in which the index of the nearest neighbors will be saved.
///
/// * `distances` - vector storing relative distances of the nearest neighbors.
///
/// * `metric_f` - metric function.
pub fn search<F>(
&self,
target: &U,
target_index: usize,
k: usize,
neighbors_indices: &mut [usize],
distances: &mut [T],
metric_f: F,
) where
F: Fn(&U, &U) -> T,
{
debug_assert_eq!(neighbors_indices.len(), distances.len());
// Use a priority queue to store intermediate results on.
let mut heap: BinaryHeap<HeapItem<T>> = BinaryHeap::with_capacity(k);
// Perform the search.
self.look_up(&mut T::max_value(), target, k, &mut heap, metric_f);
// Gather final results.
let results = heap.into_sorted_vec();
// Avoid the target itself.
neighbors_indices
.iter_mut()
.zip(distances.iter_mut())
.zip(results.iter().filter(|result| {
let HeapItem { index, distance: _ } = result;
target_index != *index
}))
.for_each(|((idx, d), result)| {
let HeapItem { index, distance } = result;
*idx = *index;
*d = *distance;
});
}
}