94 lines
2.7 KiB
Rust
94 lines
2.7 KiB
Rust
extern crate num;
|
|
|
|
use std::fmt;
|
|
use std::ops::Index;
|
|
use std::cmp::PartialOrd;
|
|
|
|
pub trait Zero {
|
|
fn zero() -> Self;
|
|
}
|
|
|
|
impl<T: num::Zero> Zero for T {
|
|
fn zero() -> Self {
|
|
num::Zero::zero()
|
|
}
|
|
}
|
|
|
|
/// 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, U> KdTree<T> where T: Copy + Clone + Index<usize, Output = U> + Zero, U: PartialOrd {
|
|
pub fn new(points: &mut [T], dimension: usize, leaf_size: usize) -> KdTree<T> {
|
|
KdTree::new_with_depth(points, dimension, leaf_size, 0)
|
|
}
|
|
|
|
fn new_with_depth(points: &mut [T], dimension: usize, leaf_size: usize, depth: usize) -> KdTree<T> {
|
|
|
|
// Select axis to work on
|
|
let axis = depth % dimension;
|
|
|
|
// Sort the points by axis
|
|
points.sort_by(|p1, p2| p1[axis].partial_cmp(&p2[axis]).unwrap());
|
|
|
|
// 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![Zero::zero(); len];
|
|
contained.copy_from_slice(points);
|
|
return KdTree::Leaf(contained);
|
|
}
|
|
|
|
// Else split what's remaining
|
|
let delimiter = points.len() / 2;
|
|
let median = points[delimiter];
|
|
|
|
// Build the left and right branches
|
|
let zero = Zero::zero();
|
|
let mut left = vec![zero; delimiter];
|
|
left.copy_from_slice(&points[0..delimiter]);
|
|
|
|
let right = &mut points[delimiter + 1 .. len];
|
|
|
|
// Recursive call
|
|
KdTree::Node {
|
|
location: median,
|
|
left: Box::new(KdTree::new_with_depth(&mut left, dimension, leaf_size, depth + 1)),
|
|
right: Box::new(KdTree::new_with_depth(right, dimension, leaf_size, depth + 1)),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<T: Zero + Copy + Clone + fmt::Display + Index<usize>> KdTree<T> {
|
|
fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result {
|
|
match self {
|
|
&KdTree::Node { 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: Zero + Copy + Clone + fmt::Display + Index<usize>> fmt::Display for KdTree<T> {
|
|
fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
|
|
self.print(formatter, "")
|
|
}
|
|
|
|
}
|