217 lines
5.9 KiB
Rust
217 lines
5.9 KiB
Rust
use std::cmp::Ordering;
|
|
use std::fmt;
|
|
|
|
pub enum Side {
|
|
Left,
|
|
Right,
|
|
}
|
|
|
|
/// A k-dimension kd tree
|
|
pub enum KdTree<T> {
|
|
Node {
|
|
location: T,
|
|
left: Box<KdTree<T>>,
|
|
right: Box<KdTree<T>>,
|
|
},
|
|
Leaf(Vec<T>),
|
|
}
|
|
|
|
impl<T: Clone> KdTree<T> {
|
|
pub fn min_leaf_size<F: Fn(usize, &T, &T) -> Ordering>(
|
|
points: &mut [T],
|
|
dimension: usize,
|
|
comparator: F,
|
|
leaf_size: usize,
|
|
) -> KdTree<T> {
|
|
// Compute bounding box of points
|
|
KdTree::min_leaf_size_with_depth(points, dimension, &comparator, leaf_size, 0)
|
|
}
|
|
|
|
fn min_leaf_size_with_depth<F: Fn(usize, &T, &T) -> Ordering>(
|
|
points: &mut [T],
|
|
dimension: usize,
|
|
comparator: &F,
|
|
leaf_size: usize,
|
|
depth: usize,
|
|
) -> KdTree<T> {
|
|
// Select axis to work on
|
|
let axis = depth % dimension;
|
|
|
|
// Sort the points by axis
|
|
points.sort_by(|a, b| comparator(axis, a, b));
|
|
|
|
// Find the median
|
|
let len = points.len();
|
|
|
|
// If there are not enough elements for a leaf, make a leaf
|
|
if len <= leaf_size {
|
|
let mut contained = vec![];
|
|
for p in points {
|
|
contained.push(p.clone());
|
|
}
|
|
return KdTree::Leaf(contained);
|
|
}
|
|
|
|
// Else split what's remaining
|
|
let delimiter = points.len() / 2;
|
|
let median = points[delimiter].clone();
|
|
|
|
// Build the left and right branches
|
|
let mut left = vec![];
|
|
for i in &points[0..delimiter + 1] {
|
|
left.push(i.clone());
|
|
}
|
|
|
|
let right = &mut points[delimiter + 1..len];
|
|
|
|
// Recursive call
|
|
KdTree::Node {
|
|
location: median,
|
|
left: Box::new(KdTree::min_leaf_size_with_depth(
|
|
&mut left,
|
|
dimension,
|
|
comparator,
|
|
leaf_size,
|
|
depth + 1,
|
|
)),
|
|
right: Box::new(KdTree::min_leaf_size_with_depth(
|
|
right,
|
|
dimension,
|
|
comparator,
|
|
leaf_size,
|
|
depth + 1,
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub fn max_depth<F: Fn(usize, &T, &T) -> Ordering>(
|
|
points: &mut [T],
|
|
dimension: usize,
|
|
comparator: F,
|
|
max_depth: usize,
|
|
) -> KdTree<T> {
|
|
// Compute bounding box of points
|
|
KdTree::max_depth_with_depth(points, dimension, &comparator, max_depth, 0)
|
|
}
|
|
|
|
fn max_depth_with_depth<F: Fn(usize, &T, &T) -> Ordering>(
|
|
points: &mut [T],
|
|
dimension: usize,
|
|
comparator: &F,
|
|
max_depth: usize,
|
|
depth: usize,
|
|
) -> KdTree<T> {
|
|
// Select axis to work on
|
|
let axis = depth % dimension;
|
|
|
|
// Sort the points by axis
|
|
points.sort_by(|a, b| comparator(axis, a, b));
|
|
|
|
// Find the median
|
|
let len = points.len();
|
|
|
|
// If we reached max depth, make a leaf
|
|
if depth >= max_depth {
|
|
let mut contained = vec![];
|
|
for p in points {
|
|
contained.push(p.clone());
|
|
}
|
|
return KdTree::Leaf(contained);
|
|
}
|
|
|
|
// Else split what's remaining
|
|
let delimiter = points.len() / 2;
|
|
let median = points[delimiter].clone();
|
|
|
|
// Build the left and right branches
|
|
let mut left = vec![];
|
|
for i in &points[0..delimiter + 1] {
|
|
left.push(i.clone());
|
|
}
|
|
|
|
let right = &mut points[delimiter + 1..len];
|
|
|
|
// Recursive call
|
|
KdTree::Node {
|
|
location: median,
|
|
left: Box::new(KdTree::max_depth_with_depth(
|
|
&mut left,
|
|
dimension,
|
|
comparator,
|
|
max_depth,
|
|
depth + 1,
|
|
)),
|
|
right: Box::new(KdTree::max_depth_with_depth(
|
|
right,
|
|
dimension,
|
|
comparator,
|
|
max_depth,
|
|
depth + 1,
|
|
)),
|
|
}
|
|
}
|
|
|
|
// Fn(location, elements)
|
|
pub fn traverse_leaves<F: FnMut(&Vec<(T, Side)>, &Vec<T>)>(&self, callback: &mut F) {
|
|
self.traverse_leaves_aux(callback, &mut vec![])
|
|
}
|
|
|
|
pub fn traverse_leaves_aux<F: FnMut(&Vec<(T, Side)>, &Vec<T>)>(
|
|
&self,
|
|
callback: &mut F,
|
|
locations: &mut Vec<(T, Side)>,
|
|
) {
|
|
// Compute the full bounding box of the points
|
|
match *self {
|
|
KdTree::Node {
|
|
ref location,
|
|
ref left,
|
|
ref right,
|
|
} => {
|
|
locations.push((location.clone(), Side::Left));
|
|
left.traverse_leaves_aux(callback, locations);
|
|
locations.pop();
|
|
|
|
locations.push((location.clone(), Side::Right));
|
|
right.traverse_leaves_aux(callback, locations);
|
|
locations.pop();
|
|
}
|
|
KdTree::Leaf(ref elements) => {
|
|
callback(locations, elements);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<T: fmt::Display> KdTree<T> {
|
|
fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result {
|
|
match self {
|
|
&KdTree::Node {
|
|
ref location,
|
|
ref left,
|
|
ref right,
|
|
} => {
|
|
writeln!(formatter, "{}{}", indent, location)?;
|
|
left.print(formatter, &format!("{} ", indent))?;
|
|
right.print(formatter, &format!("{} ", indent))?;
|
|
}
|
|
|
|
&KdTree::Leaf(ref elements) => {
|
|
write!(formatter, "{}[", indent)?;
|
|
for element in elements {
|
|
write!(formatter, "{}, ", element)?;
|
|
}
|
|
writeln!(formatter, "]")?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl<T: fmt::Display> fmt::Display for KdTree<T> {
|
|
fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
|
|
self.print(formatter, "")
|
|
}
|
|
}
|