Create kd tree by depth
This commit is contained in:
parent
2dc25e95f8
commit
da46959258
|
@ -31,7 +31,7 @@ fn main() {
|
|||
Vector2::new(7.0, 2.0),
|
||||
];
|
||||
|
||||
let kdtree = KdTree::new(
|
||||
let kdtree = KdTree::min_leaf_size(
|
||||
&mut points,
|
||||
2,
|
||||
|coord, ref a, ref b| a.data[coord].partial_cmp(&b.data[coord]).unwrap(),
|
||||
|
|
105
src/lib.rs
105
src/lib.rs
|
@ -1,5 +1,5 @@
|
|||
use std::fmt;
|
||||
use std::cmp::Ordering;
|
||||
use std::fmt;
|
||||
|
||||
pub enum Side {
|
||||
Left,
|
||||
|
@ -16,18 +16,18 @@ pub enum KdTree<T> {
|
|||
Leaf(Vec<T>),
|
||||
}
|
||||
|
||||
impl<T: Clone> KdTree<T> where {
|
||||
pub fn new<F: Fn(usize, &T, &T) -> Ordering>(
|
||||
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::new_with_depth(points, dimension, &comparator, leaf_size, 0)
|
||||
KdTree::min_leaf_size_with_depth(points, dimension, &comparator, leaf_size, 0)
|
||||
}
|
||||
|
||||
fn new_with_depth<F: Fn(usize, &T, &T) -> Ordering>(
|
||||
fn min_leaf_size_with_depth<F: Fn(usize, &T, &T) -> Ordering>(
|
||||
points: &mut [T],
|
||||
dimension: usize,
|
||||
comparator: &F,
|
||||
|
@ -67,14 +67,14 @@ impl<T: Clone> KdTree<T> where {
|
|||
// Recursive call
|
||||
KdTree::Node {
|
||||
location: median,
|
||||
left: Box::new(KdTree::new_with_depth(
|
||||
left: Box::new(KdTree::min_leaf_size_with_depth(
|
||||
&mut left,
|
||||
dimension,
|
||||
comparator,
|
||||
leaf_size,
|
||||
depth + 1,
|
||||
)),
|
||||
right: Box::new(KdTree::new_with_depth(
|
||||
right: Box::new(KdTree::min_leaf_size_with_depth(
|
||||
right,
|
||||
dimension,
|
||||
comparator,
|
||||
|
@ -84,17 +84,90 @@ impl<T: Clone> KdTree<T> where {
|
|||
}
|
||||
}
|
||||
|
||||
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)>) {
|
||||
|
||||
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 } => {
|
||||
|
||||
KdTree::Node {
|
||||
ref location,
|
||||
ref left,
|
||||
ref right,
|
||||
} => {
|
||||
locations.push((location.clone(), Side::Left));
|
||||
left.traverse_leaves_aux(callback, locations);
|
||||
locations.pop();
|
||||
|
@ -102,12 +175,10 @@ impl<T: Clone> KdTree<T> where {
|
|||
locations.push((location.clone(), Side::Right));
|
||||
right.traverse_leaves_aux(callback, locations);
|
||||
locations.pop();
|
||||
|
||||
},
|
||||
}
|
||||
KdTree::Leaf(ref elements) => {
|
||||
callback(locations, elements);
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
|
@ -115,7 +186,11 @@ impl<T: Clone> KdTree<T> where {
|
|||
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, } => {
|
||||
&KdTree::Node {
|
||||
ref location,
|
||||
ref left,
|
||||
ref right,
|
||||
} => {
|
||||
writeln!(formatter, "{}{}", indent, location)?;
|
||||
left.print(formatter, &format!("{} ", indent))?;
|
||||
right.print(formatter, &format!("{} ", indent))?;
|
||||
|
|
Loading…
Reference in New Issue