kiddo 5.0.3

A high-performance, flexible, ergonomic k-d tree library. Ideal for geo- and astro- nearest-neighbour and k-nearest-neighbor queries
Documentation 
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

use crate::float::kdtree::Axis;
use crate::nearest_neighbour::NearestNeighbour;
use crate::traits::Content;
use sorted_vec::SortedVec;
use std::collections::BinaryHeap;

pub trait ResultCollection<A: Axis, T: Content> {
    fn new_with_capacity(capacity: usize) -> Self;
    fn add(&mut self, entry: NearestNeighbour<A, T>);
    fn max_dist(&self) -> A;
    fn into_vec(self) -> Vec<NearestNeighbour<A, T>>;
    fn into_sorted_vec(self) -> Vec<NearestNeighbour<A, T>>;
}

impl<A: Axis, T: Content> ResultCollection<A, T> for BinaryHeap<NearestNeighbour<A, T>> {
    fn new_with_capacity(capacity: usize) -> Self {
        BinaryHeap::with_capacity(capacity)
    }
    fn add(&mut self, entry: NearestNeighbour<A, T>) {
        let k = self.capacity();
        if self.len() < k {
            self.push(entry);
        } else {
            let mut max_heap_value = self.peek_mut().unwrap();
            if entry < *max_heap_value {
                *max_heap_value = entry;
            }
        }
    }
    fn max_dist(&self) -> A {
        if self.len() < self.capacity() {
            A::infinity()
        } else {
            self.peek().map_or(A::infinity(), |n| n.distance)
        }
    }
    fn into_vec(self) -> Vec<NearestNeighbour<A, T>> {
        BinaryHeap::into_vec(self)
    }
    fn into_sorted_vec(self) -> Vec<NearestNeighbour<A, T>> {
        BinaryHeap::into_sorted_vec(self)
    }
}

impl<A: Axis, T: Content> ResultCollection<A, T> for Vec<NearestNeighbour<A, T>> {
    fn new_with_capacity(capacity: usize) -> Self {
        Vec::with_capacity(capacity)
    }

    fn add(&mut self, entry: NearestNeighbour<A, T>) {
        self.push(entry)
    }

    fn max_dist(&self) -> A {
        A::infinity()
    }

    fn into_vec(self) -> Vec<NearestNeighbour<A, T>> {
        self
    }

    fn into_sorted_vec(mut self) -> Vec<NearestNeighbour<A, T>> {
        self.sort();
        self
    }
}

impl<A: Axis, T: Content> ResultCollection<A, T> for SortedVec<NearestNeighbour<A, T>> {
    fn new_with_capacity(capacity: usize) -> Self {
        SortedVec::with_capacity(capacity)
    }

    fn add(&mut self, entry: NearestNeighbour<A, T>) {
        let len = self.len();
        if len < self.capacity() {
            self.insert(entry);
        } else if entry < *self.last().unwrap() {
            self.pop();
            self.push(entry);
        }
    }

    fn max_dist(&self) -> A {
        if self.len() < self.capacity() {
            A::infinity()
        } else {
            self.last().map_or(A::infinity(), |n| n.distance)
        }
    }

    fn into_vec(self) -> Vec<NearestNeighbour<A, T>> {
        self.into_vec()
    }

    fn into_sorted_vec(self) -> Vec<NearestNeighbour<A, T>> {
        self.into_vec()
    }
}