87 lines
2.6 KiB
Rust
87 lines
2.6 KiB
Rust
extern crate num;
|
|
extern crate model_converter;
|
|
|
|
use std::fmt;
|
|
|
|
use num::Zero;
|
|
|
|
use model_converter::math::vector;
|
|
|
|
macro_rules! define_kdtree {
|
|
|
|
($name: ident, $vector_type: ty, $dimension: expr) => {
|
|
|
|
/// A 3-dimension kd tree
|
|
pub struct $name {
|
|
location: $vector_type,
|
|
left: Option<Box<$name>>,
|
|
right: Option<Box<$name>>,
|
|
}
|
|
|
|
impl $name {
|
|
pub fn new(points: &mut [$vector_type]) -> $name {
|
|
$name::new_with_depth(points, 0)
|
|
}
|
|
|
|
fn new_with_depth(points: &mut [$vector_type], depth: usize) -> $name {
|
|
|
|
// 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();
|
|
let delimiter = points.len() / 2;
|
|
let median = points[delimiter];
|
|
|
|
let zero = Zero::zero();
|
|
let mut left = vec![zero; delimiter];
|
|
left.copy_from_slice(&points[0..delimiter]);
|
|
|
|
let right = &mut points[delimiter + 1 .. len];
|
|
|
|
if left.len() == 0 || right.len() == 0 {
|
|
$name {
|
|
location: median,
|
|
left: None,
|
|
right: None,
|
|
}
|
|
} else {
|
|
$name {
|
|
location: median,
|
|
left: Some(Box::new($name::new_with_depth(&mut left, depth + 1))),
|
|
right: Some(Box::new($name::new_with_depth(right, depth + 1))),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result {
|
|
writeln!(formatter, "{}{}", indent, self.location)?;
|
|
|
|
if let Some(left) = self.left.as_ref() {
|
|
left.print(formatter, &format!("{} ", indent))?;
|
|
}
|
|
|
|
if let Some(right) = self.right.as_ref() {
|
|
right.print(formatter, &format!("{} ", indent))?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl fmt::Display for $name {
|
|
fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
|
|
self.print(formatter, "")
|
|
}
|
|
|
|
}
|
|
}
|
|
}
|
|
|
|
define_kdtree!(KdTree2, vector::Vector2<f64>, 2);
|
|
define_kdtree!(KdTree3, vector::Vector3<f64>, 3);
|
|
define_kdtree!(KdTree4, vector::Vector4<f64>, 4);
|