summaryrefslogtreecommitdiffstats
path: root/src/kd.rs
diff options
context:
space:
mode:
authorTavian Barnes <tavianator@tavianator.com>2020-05-27 14:49:32 -0400
committerTavian Barnes <tavianator@tavianator.com>2020-06-24 10:02:23 -0400
commitaf325f5a4ff967314fb484b77b33406d21133393 (patch)
treee7a45c65708c7b8bc5a8c0b6230d713fe11950e2 /src/kd.rs
parenta6e586d3f1368e4ecbe56b6481e8bca2ec8e7bb9 (diff)
downloadacap-af325f5a4ff967314fb484b77b33406d21133393.tar.xz
kd: Implement flat k-d trees
Diffstat (limited to 'src/kd.rs')
-rw-r--r--src/kd.rs155
1 files changed, 155 insertions, 0 deletions
diff --git a/src/kd.rs b/src/kd.rs
index 97616e7..4f591f9 100644
--- a/src/kd.rs
+++ b/src/kd.rs
@@ -333,6 +333,156 @@ where
V: Coordinates,
{}
+/// A node in a flat k-d tree.
+#[derive(Debug)]
+struct FlatKdNode<T> {
+ /// The vantage point itself.
+ item: T,
+ /// The size of the left subtree.
+ left_len: usize,
+}
+
+impl<T: Coordinates> FlatKdNode<T> {
+ /// Create a new FlatKdNode.
+ fn new(item: T) -> Self {
+ Self {
+ item,
+ left_len: 0,
+ }
+ }
+
+ /// Create a balanced tree.
+ fn balanced<I: IntoIterator<Item = T>>(items: I) -> Vec<Self> {
+ let mut nodes: Vec<_> = items
+ .into_iter()
+ .map(Self::new)
+ .collect();
+
+ Self::balance_recursive(&mut nodes, 0);
+
+ nodes
+ }
+
+ /// Create a balanced subtree.
+ fn balance_recursive(nodes: &mut [Self], level: usize) {
+ if !nodes.is_empty() {
+ nodes.sort_by_cached_key(|x| Ordered::new(x.item.coord(level)));
+
+ let mid = nodes.len() / 2;
+ nodes.swap(0, mid);
+
+ let (node, children) = nodes.split_first_mut().unwrap();
+ let (left, right) = children.split_at_mut(mid);
+ node.left_len = left.len();
+
+ let next = (level + 1) % node.item.dims();
+ Self::balance_recursive(left, next);
+ Self::balance_recursive(right, next);
+ }
+ }
+}
+
+impl<'a, K, V, N> KdSearch<K, &'a V, N> for &'a [FlatKdNode<V>]
+where
+ K: KdProximity<&'a V>,
+ V: Coordinates,
+ N: Neighborhood<K, &'a V>,
+{
+ fn item(self) -> &'a V {
+ &self[0].item
+ }
+
+ fn left(self) -> Option<Self> {
+ let end = self[0].left_len + 1;
+ if end > 1 {
+ Some(&self[1..end])
+ } else {
+ None
+ }
+ }
+
+ fn right(self) -> Option<Self> {
+ let start = self[0].left_len + 1;
+ if start < self.len() {
+ Some(&self[start..])
+ } else {
+ None
+ }
+ }
+}
+
+/// A [k-d tree] stored as a flat array.
+///
+/// A FlatKdTree is always balanced and usually more efficient than a [KdTree], but doesn't support
+/// dynamic updates.
+///
+/// [k-d tree]: https://en.wikipedia.org/wiki/K-d_tree
+#[derive(Debug)]
+pub struct FlatKdTree<T> {
+ nodes: Vec<FlatKdNode<T>>,
+}
+
+impl<T: Coordinates> FlatKdTree<T> {
+ /// Create a balanced tree out of a sequence of items.
+ pub fn balanced<I: IntoIterator<Item = T>>(items: I) -> Self {
+ Self {
+ nodes: FlatKdNode::balanced(items),
+ }
+ }
+}
+
+impl<T: Coordinates> FromIterator<T> for FlatKdTree<T> {
+ fn from_iter<I: IntoIterator<Item = T>>(items: I) -> Self {
+ Self::balanced(items)
+ }
+}
+
+/// An iterator that moves values out of a flat k-d tree.
+#[derive(Debug)]
+pub struct FlatIntoIter<T>(std::vec::IntoIter<FlatKdNode<T>>);
+
+impl<T> Iterator for FlatIntoIter<T> {
+ type Item = T;
+
+ fn next(&mut self) -> Option<T> {
+ self.0.next().map(|n| n.item)
+ }
+}
+
+impl<T> IntoIterator for FlatKdTree<T> {
+ type Item = T;
+ type IntoIter = FlatIntoIter<T>;
+
+ fn into_iter(self) -> Self::IntoIter {
+ FlatIntoIter(self.nodes.into_iter())
+ }
+}
+
+impl<K, V> NearestNeighbors<K, V> for FlatKdTree<V>
+where
+ K: KdProximity<V>,
+ V: Coordinates,
+{
+ fn search<'k, 'v, N>(&'v self, mut neighborhood: N) -> N
+ where
+ K: 'k,
+ V: 'v,
+ N: Neighborhood<&'k K, &'v V>,
+ {
+ if !self.nodes.is_empty() {
+ let mut closest = neighborhood.target().as_vec();
+ self.nodes.as_slice().search(0, &mut closest, &mut neighborhood);
+ }
+ neighborhood
+ }
+}
+
+impl<K, V> ExactNeighbors<K, V> for FlatKdTree<V>
+where
+ K: KdMetric<V>,
+ V: Coordinates,
+{}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -354,4 +504,9 @@ mod tests {
tree
});
}
+
+ #[test]
+ fn test_flat_kd_tree() {
+ test_nearest_neighbors(FlatKdTree::from_iter);
+ }
}