Skip to main content

max / audiofiles

15.2 KB · 522 lines History Blame Raw
1 //! Vantage-point tree for sub-linear nearest-neighbor and range queries in metric spaces.
2 //!
3 //! A VP-tree partitions data by distance from selected "vantage points", enabling
4 //! efficient pruning during search. Requires a distance function satisfying the
5 //! triangle inequality.
6 //!
7 //! - Build: O(n log n) distance computations
8 //! - k-nearest query: O(log n) expected distance computations
9 //! - Range query: O(log n + results) expected distance computations
10
11 use std::collections::BinaryHeap;
12
13 /// A vantage-point tree for nearest-neighbor and range queries.
14 pub struct VpTree<T> {
15 items: Vec<T>,
16 root: Option<usize>,
17 nodes: Vec<VpNode>,
18 }
19
20 struct VpNode {
21 /// Index into `items`.
22 item_idx: usize,
23 /// Median distance from this vantage point to items in its subtree.
24 threshold: f64,
25 /// Inside subtree: items with distance <= threshold.
26 left: Option<usize>,
27 /// Outside subtree: items with distance > threshold.
28 right: Option<usize>,
29 }
30
31 /// A search result: index into the original items vector plus distance from the query.
32 #[derive(Debug, Clone)]
33 pub struct VpMatch {
34 pub index: usize,
35 pub distance: f64,
36 }
37
38 impl<T> VpTree<T> {
39 /// Build a VP-tree from a set of items.
40 ///
41 /// The distance function must be a metric (non-negative, symmetric, triangle
42 /// inequality). Build cost is O(n log n) distance computations.
43 pub fn build(items: Vec<T>, dist: impl Fn(&T, &T) -> f64) -> Self {
44 let n = items.len();
45 if n == 0 {
46 return Self {
47 items,
48 root: None,
49 nodes: Vec::new(),
50 };
51 }
52 let mut indices: Vec<usize> = (0..n).collect();
53 let mut nodes = Vec::with_capacity(n);
54 let root = build_recursive(&items, &mut indices, &dist, &mut nodes, 0);
55 Self {
56 items,
57 root: Some(root),
58 nodes,
59 }
60 }
61
62 /// Find the `k` nearest items to `query`, sorted by ascending distance.
63 pub fn find_nearest(
64 &self,
65 query: &T,
66 k: usize,
67 dist: impl Fn(&T, &T) -> f64,
68 ) -> Vec<VpMatch> {
69 if self.items.is_empty() || k == 0 {
70 return Vec::new();
71 }
72 let mut heap: BinaryHeap<HeapEntry> = BinaryHeap::with_capacity(k + 1);
73 let mut tau = f64::INFINITY;
74 if let Some(root) = self.root {
75 search_nearest(
76 &self.items,
77 &self.nodes,
78 root,
79 query,
80 k,
81 &dist,
82 &mut heap,
83 &mut tau,
84 );
85 }
86 let mut results: Vec<VpMatch> = heap
87 .into_iter()
88 .map(|e| VpMatch {
89 index: e.index,
90 distance: e.distance,
91 })
92 .collect();
93 results.sort_by(|a, b| a.distance.total_cmp(&b.distance));
94 results
95 }
96
97 /// Find all items within `radius` of `query`, sorted by ascending distance.
98 pub fn find_within(
99 &self,
100 query: &T,
101 radius: f64,
102 dist: impl Fn(&T, &T) -> f64,
103 ) -> Vec<VpMatch> {
104 if self.items.is_empty() {
105 return Vec::new();
106 }
107 let mut results = Vec::new();
108 if let Some(root) = self.root {
109 search_within(
110 &self.items,
111 &self.nodes,
112 root,
113 query,
114 radius,
115 &dist,
116 &mut results,
117 );
118 }
119 results.sort_by(|a, b| a.distance.total_cmp(&b.distance));
120 results
121 }
122
123 /// Number of items in the tree.
124 pub fn len(&self) -> usize {
125 self.items.len()
126 }
127
128 /// Whether the tree is empty.
129 pub fn is_empty(&self) -> bool {
130 self.items.is_empty()
131 }
132
133 /// Reference to an item by index.
134 pub fn get(&self, index: usize) -> &T {
135 &self.items[index]
136 }
137 }
138
139 // --- Build ---
140
141 /// Maximum recursion depth to prevent stack overflow on degenerate input
142 /// (e.g. all-identical feature vectors causing O(n)-deep recursion).
143 const MAX_BUILD_DEPTH: usize = 64;
144
145 fn build_recursive<T>(
146 items: &[T],
147 indices: &mut [usize],
148 dist: &impl Fn(&T, &T) -> f64,
149 nodes: &mut Vec<VpNode>,
150 depth: usize,
151 ) -> usize {
152 debug_assert!(!indices.is_empty());
153
154 let vp_idx = indices[0];
155 // Leaf node: single item
156 if indices.len() == 1 {
157 let node_idx = nodes.len();
158 nodes.push(VpNode {
159 item_idx: vp_idx,
160 threshold: 0.0,
161 left: None,
162 right: None,
163 });
164 return node_idx;
165 }
166
167 // Depth cap: chain remaining items as a flat linked list (iteratively)
168 // so none are lost. Each node holds one item, threshold=INFINITY ensures
169 // the search always visits the left child.
170 if depth >= MAX_BUILD_DEPTH {
171 let first_node_idx = nodes.len();
172 // Create all leaf nodes
173 for &idx in indices.iter() {
174 nodes.push(VpNode {
175 item_idx: idx,
176 threshold: f64::INFINITY,
177 left: None,
178 right: None,
179 });
180 }
181 // Link them: each node's left points to the next
182 for i in 0..indices.len() - 1 {
183 nodes[first_node_idx + i].left = Some(first_node_idx + i + 1);
184 }
185 return first_node_idx;
186 }
187
188 // Compute distances from vantage point to all other items in this subtree.
189 let mut dists: Vec<(usize, f64)> = indices[1..]
190 .iter()
191 .map(|&idx| (idx, dist(&items[vp_idx], &items[idx])))
192 .collect();
193
194 // Partition at median distance.
195 let median_pos = dists.len() / 2;
196 dists.select_nth_unstable_by(median_pos, |a, b| a.1.total_cmp(&b.1));
197 let threshold = dists[median_pos].1;
198
199 // Split into inside (<= threshold) and outside (> threshold).
200 let mut inside = Vec::with_capacity(median_pos + 1);
201 let mut outside = Vec::with_capacity(dists.len() - median_pos);
202 for &(idx, d) in &dists {
203 if d <= threshold {
204 inside.push(idx);
205 } else {
206 outside.push(idx);
207 }
208 }
209
210 // Allocate node, fill children after recursion.
211 let node_idx = nodes.len();
212 nodes.push(VpNode {
213 item_idx: vp_idx,
214 threshold,
215 left: None,
216 right: None,
217 });
218
219 let left = if inside.is_empty() {
220 None
221 } else {
222 Some(build_recursive(items, &mut inside, dist, nodes, depth + 1))
223 };
224 let right = if outside.is_empty() {
225 None
226 } else {
227 Some(build_recursive(items, &mut outside, dist, nodes, depth + 1))
228 };
229
230 nodes[node_idx].left = left;
231 nodes[node_idx].right = right;
232 node_idx
233 }
234
235 // --- k-nearest search ---
236
237 struct HeapEntry {
238 distance: f64,
239 index: usize,
240 }
241
242 impl Eq for HeapEntry {}
243 impl PartialEq for HeapEntry {
244 fn eq(&self, other: &Self) -> bool {
245 self.distance.total_cmp(&other.distance) == std::cmp::Ordering::Equal
246 }
247 }
248 impl PartialOrd for HeapEntry {
249 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
250 Some(self.cmp(other))
251 }
252 }
253 impl Ord for HeapEntry {
254 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
255 self.distance.total_cmp(&other.distance)
256 }
257 }
258
259 #[allow(clippy::too_many_arguments)]
260 fn search_nearest<T>(
261 items: &[T],
262 nodes: &[VpNode],
263 node_idx: usize,
264 query: &T,
265 k: usize,
266 dist: &impl Fn(&T, &T) -> f64,
267 heap: &mut BinaryHeap<HeapEntry>,
268 tau: &mut f64,
269 ) {
270 let node = &nodes[node_idx];
271 let d = dist(query, &items[node.item_idx]);
272
273 // Consider this node's vantage point.
274 if d < *tau || heap.len() < k {
275 heap.push(HeapEntry {
276 distance: d,
277 index: node.item_idx,
278 });
279 if heap.len() > k {
280 heap.pop();
281 }
282 if heap.len() == k {
283 *tau = heap.peek().unwrap().distance;
284 }
285 }
286
287 // Search the closer subtree first for better pruning.
288 if d <= node.threshold {
289 // Query is inside — search inside first.
290 if let Some(left) = node.left {
291 if d - *tau <= node.threshold {
292 search_nearest(items, nodes, left, query, k, dist, heap, tau);
293 }
294 }
295 if let Some(right) = node.right {
296 if d + *tau > node.threshold {
297 search_nearest(items, nodes, right, query, k, dist, heap, tau);
298 }
299 }
300 } else {
301 // Query is outside — search outside first.
302 if let Some(right) = node.right {
303 if d + *tau > node.threshold {
304 search_nearest(items, nodes, right, query, k, dist, heap, tau);
305 }
306 }
307 if let Some(left) = node.left {
308 if d - *tau <= node.threshold {
309 search_nearest(items, nodes, left, query, k, dist, heap, tau);
310 }
311 }
312 }
313 }
314
315 // --- Range search ---
316
317 fn search_within<T>(
318 items: &[T],
319 nodes: &[VpNode],
320 node_idx: usize,
321 query: &T,
322 radius: f64,
323 dist: &impl Fn(&T, &T) -> f64,
324 results: &mut Vec<VpMatch>,
325 ) {
326 let node = &nodes[node_idx];
327 let d = dist(query, &items[node.item_idx]);
328
329 if d <= radius {
330 results.push(VpMatch {
331 index: node.item_idx,
332 distance: d,
333 });
334 }
335
336 // Prune subtrees using triangle inequality bounds.
337 if let Some(left) = node.left {
338 if d - radius <= node.threshold {
339 search_within(items, nodes, left, query, radius, dist, results);
340 }
341 }
342 if let Some(right) = node.right {
343 if d + radius > node.threshold {
344 search_within(items, nodes, right, query, radius, dist, results);
345 }
346 }
347 }
348
349 #[cfg(test)]
350 mod tests {
351 use super::*;
352
353 fn euclidean_1d(a: &f64, b: &f64) -> f64 {
354 (a - b).abs()
355 }
356
357 fn euclidean_2d(a: &[f64; 2], b: &[f64; 2]) -> f64 {
358 ((a[0] - b[0]).powi(2) + (a[1] - b[1]).powi(2)).sqrt()
359 }
360
361 #[test]
362 fn empty_tree() {
363 let tree: VpTree<f64> = VpTree::build(vec![], euclidean_1d);
364 assert!(tree.is_empty());
365 assert_eq!(tree.len(), 0);
366 assert!(tree.find_nearest(&0.0, 5, euclidean_1d).is_empty());
367 assert!(tree.find_within(&0.0, 1.0, euclidean_1d).is_empty());
368 }
369
370 #[test]
371 fn single_item() {
372 let tree = VpTree::build(vec![5.0], euclidean_1d);
373 assert_eq!(tree.len(), 1);
374 let results = tree.find_nearest(&5.0, 1, euclidean_1d);
375 assert_eq!(results.len(), 1);
376 assert!(results[0].distance.abs() < f64::EPSILON);
377 }
378
379 #[test]
380 fn k_nearest_basic() {
381 let items = vec![1.0, 3.0, 5.0, 7.0, 9.0, 11.0, 13.0, 15.0];
382 let tree = VpTree::build(items, euclidean_1d);
383 let results = tree.find_nearest(&6.0, 3, euclidean_1d);
384 assert_eq!(results.len(), 3);
385 let values: Vec<f64> = results.iter().map(|r| *tree.get(r.index)).collect();
386 assert!(values.contains(&5.0));
387 assert!(values.contains(&7.0));
388 }
389
390 #[test]
391 fn find_within_basic() {
392 let items = vec![1.0, 3.0, 5.0, 7.0, 9.0];
393 let tree = VpTree::build(items, euclidean_1d);
394 let results = tree.find_within(&5.0, 2.5, euclidean_1d);
395 let values: Vec<f64> = results.iter().map(|r| *tree.get(r.index)).collect();
396 assert!(values.contains(&3.0));
397 assert!(values.contains(&5.0));
398 assert!(values.contains(&7.0));
399 assert_eq!(values.len(), 3);
400 }
401
402 #[test]
403 fn results_sorted_by_distance() {
404 let items: Vec<f64> = (0..100).map(|i| i as f64).collect();
405 let tree = VpTree::build(items, euclidean_1d);
406
407 let nearest = tree.find_nearest(&50.0, 10, euclidean_1d);
408 for w in nearest.windows(2) {
409 assert!(w[0].distance <= w[1].distance);
410 }
411
412 let within = tree.find_within(&50.0, 5.0, euclidean_1d);
413 for w in within.windows(2) {
414 assert!(w[0].distance <= w[1].distance);
415 }
416 }
417
418 #[test]
419 fn k_larger_than_n() {
420 let items = vec![1.0, 2.0, 3.0];
421 let tree = VpTree::build(items, euclidean_1d);
422 let results = tree.find_nearest(&0.0, 10, euclidean_1d);
423 assert_eq!(results.len(), 3);
424 }
425
426 #[test]
427 fn k_zero() {
428 let items = vec![1.0, 2.0, 3.0];
429 let tree = VpTree::build(items, euclidean_1d);
430 let results = tree.find_nearest(&0.0, 0, euclidean_1d);
431 assert!(results.is_empty());
432 }
433
434 #[test]
435 fn two_dimensional() {
436 let items = vec![
437 [0.0, 0.0],
438 [1.0, 0.0],
439 [0.0, 1.0],
440 [1.0, 1.0],
441 [10.0, 10.0],
442 ];
443 let tree = VpTree::build(items, euclidean_2d);
444 let results = tree.find_nearest(&[0.5, 0.5], 4, euclidean_2d);
445 assert_eq!(results.len(), 4);
446 // [10, 10] should not be in top-4
447 assert!(!results.iter().any(|r| tree.get(r.index)[0] > 5.0));
448 }
449
450 #[test]
451 fn find_within_excludes_far_items() {
452 let items = vec![0.0, 1.0, 2.0, 100.0, 200.0];
453 let tree = VpTree::build(items, euclidean_1d);
454 let results = tree.find_within(&1.0, 1.5, euclidean_1d);
455 assert_eq!(results.len(), 3); // 0.0, 1.0, 2.0
456 assert!(results.iter().all(|r| *tree.get(r.index) <= 2.5));
457 }
458
459 #[test]
460 fn correctness_vs_brute_force() {
461 // 1000 deterministic pseudo-random points.
462 let items: Vec<f64> = (0..1000)
463 .map(|i| ((i as f64 * 0.618033988749895) % 1.0) * 100.0)
464 .collect();
465 let tree = VpTree::build(items.clone(), euclidean_1d);
466 let query = 42.0;
467
468 // k-nearest
469 let k = 10;
470 let tree_results = tree.find_nearest(&query, k, euclidean_1d);
471 let mut brute: Vec<(usize, f64)> = items
472 .iter()
473 .enumerate()
474 .map(|(i, &v)| (i, (v - query).abs()))
475 .collect();
476 brute.sort_by(|a, b| a.1.total_cmp(&b.1));
477
478 assert_eq!(tree_results.len(), k);
479 for (tr, br) in tree_results.iter().zip(brute.iter()) {
480 assert!(
481 (tr.distance - br.1).abs() < 1e-10,
482 "VP-tree distance {} != brute force distance {}",
483 tr.distance,
484 br.1
485 );
486 }
487
488 // find_within
489 let radius = 5.0;
490 let tree_within = tree.find_within(&query, radius, euclidean_1d);
491 let brute_within: Vec<f64> = items
492 .iter()
493 .filter(|&&v| (v - query).abs() <= radius)
494 .copied()
495 .collect();
496 assert_eq!(
497 tree_within.len(),
498 brute_within.len(),
499 "VP-tree found {} items within radius, brute force found {}",
500 tree_within.len(),
501 brute_within.len()
502 );
503 }
504
505 #[test]
506 fn two_items() {
507 let tree = VpTree::build(vec![0.0, 10.0], euclidean_1d);
508 let results = tree.find_nearest(&3.0, 1, euclidean_1d);
509 assert_eq!(results.len(), 1);
510 assert!((tree.get(results[0].index) - 0.0).abs() < f64::EPSILON);
511 }
512
513 #[test]
514 fn duplicate_distances() {
515 // Multiple items at the same distance from each other.
516 let items = vec![0.0, 5.0, 5.0, 5.0, 10.0];
517 let tree = VpTree::build(items, euclidean_1d);
518 let results = tree.find_within(&5.0, 0.0, euclidean_1d);
519 assert_eq!(results.len(), 3); // three items at distance 0
520 }
521 }
522