Create kd tree by depth

This commit is contained in:
Thomas Forgione 2022-08-16 15:40:46 +02:00
parent 2dc25e95f8
commit da46959258
2 changed files with 91 additions and 16 deletions

View File

@ -31,7 +31,7 @@ fn main() {
Vector2::new(7.0, 2.0), Vector2::new(7.0, 2.0),
]; ];
let kdtree = KdTree::new( let kdtree = KdTree::min_leaf_size(
&mut points, &mut points,
2, 2,
|coord, ref a, ref b| a.data[coord].partial_cmp(&b.data[coord]).unwrap(), |coord, ref a, ref b| a.data[coord].partial_cmp(&b.data[coord]).unwrap(),

View File

@ -1,5 +1,5 @@
use std::fmt;
use std::cmp::Ordering; use std::cmp::Ordering;
use std::fmt;
pub enum Side { pub enum Side {
Left, Left,
@ -16,18 +16,18 @@ pub enum KdTree<T> {
Leaf(Vec<T>), Leaf(Vec<T>),
} }
impl<T: Clone> KdTree<T> where { impl<T: Clone> KdTree<T> {
pub fn new<F: Fn(usize, &T, &T) -> Ordering>( pub fn min_leaf_size<F: Fn(usize, &T, &T) -> Ordering>(
points: &mut [T], points: &mut [T],
dimension: usize, dimension: usize,
comparator: F, comparator: F,
leaf_size: usize, leaf_size: usize,
) -> KdTree<T> { ) -> KdTree<T> {
// Compute bounding box of points // 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], points: &mut [T],
dimension: usize, dimension: usize,
comparator: &F, comparator: &F,
@ -67,14 +67,14 @@ impl<T: Clone> KdTree<T> where {
// Recursive call // Recursive call
KdTree::Node { KdTree::Node {
location: median, location: median,
left: Box::new(KdTree::new_with_depth( left: Box::new(KdTree::min_leaf_size_with_depth(
&mut left, &mut left,
dimension, dimension,
comparator, comparator,
leaf_size, leaf_size,
depth + 1, depth + 1,
)), )),
right: Box::new(KdTree::new_with_depth( right: Box::new(KdTree::min_leaf_size_with_depth(
right, right,
dimension, dimension,
comparator, 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) // Fn(location, elements)
pub fn traverse_leaves<F: FnMut(&Vec<(T, Side)>, &Vec<T>)>(&self, callback: &mut F) { pub fn traverse_leaves<F: FnMut(&Vec<(T, Side)>, &Vec<T>)>(&self, callback: &mut F) {
self.traverse_leaves_aux(callback, &mut vec![]) 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 // Compute the full bounding box of the points
match *self { match *self {
KdTree::Node { ref location , ref left, ref right } => { KdTree::Node {
ref location,
ref left,
ref right,
} => {
locations.push((location.clone(), Side::Left)); locations.push((location.clone(), Side::Left));
left.traverse_leaves_aux(callback, locations); left.traverse_leaves_aux(callback, locations);
locations.pop(); locations.pop();
@ -102,12 +175,10 @@ impl<T: Clone> KdTree<T> where {
locations.push((location.clone(), Side::Right)); locations.push((location.clone(), Side::Right));
right.traverse_leaves_aux(callback, locations); right.traverse_leaves_aux(callback, locations);
locations.pop(); locations.pop();
}
},
KdTree::Leaf(ref elements) => { KdTree::Leaf(ref elements) => {
callback(locations, elements); callback(locations, elements);
} }
} }
} }
} }
@ -115,7 +186,11 @@ impl<T: Clone> KdTree<T> where {
impl<T: fmt::Display> KdTree<T> { impl<T: fmt::Display> KdTree<T> {
fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result { fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result {
match self { match self {
&KdTree::Node { ref location, ref left, ref right, } => { &KdTree::Node {
ref location,
ref left,
ref right,
} => {
writeln!(formatter, "{}{}", indent, location)?; writeln!(formatter, "{}{}", indent, location)?;
left.print(formatter, &format!("{} ", indent))?; left.print(formatter, &format!("{} ", indent))?;
right.print(formatter, &format!("{} ", indent))?; right.print(formatter, &format!("{} ", indent))?;