Skip to content

Instantly share code, notes, and snippets.

@djmcgill
Created October 7, 2017 16:15
Show Gist options
  • Save djmcgill/db834d575559fb883421762503db595a to your computer and use it in GitHub Desktop.
Save djmcgill/db834d575559fb883421762503db595a to your computer and use it in GitHub Desktop.
kd_tree simplification 2
use na::*;
use na::allocator::Allocator;
/// A kd tree
pub struct KdTree<N: DimName>(Box<kd_tree::KdTreeImpl<N>>)
where
N: DimName,
DefaultAllocator: Allocator<f64, N>;
impl<N: DimName> KdTree<N>
where
N: DimName,
DefaultAllocator: Allocator<f64, N>,
{
// /// Get all the points from the kd-tree
// fn get_points(self) -> Vec<VectorN<f64, N>> {
// let mut pts = Vec::new();
// let KdTree(tree) = self;
// kd_tree::add_points_to_vector(tree, &mut pts);
// pts
// }
// fn build(pts : &Vec<VectorN<f64, N>>) -> Self {
// KdTree(Box::new(kd_tree::build(pts, 0)))
// }
}
mod kd_tree {
use na::*;
use na::allocator::Allocator;
/// kd-Tree implementation structure
pub(super) enum KdTreeImpl<N>
where
N: DimName,
DefaultAllocator: Allocator<f64, N>,
{
Node(usize, Box<KdTreeImpl<N>>, VectorN<f64, N>, Box<KdTreeImpl<N>>),
Empty(),
}
impl<N> KdTreeImpl<N>
where
N: DimName,
DefaultAllocator: Allocator<f64, N>,
{
/// mutably enunerate points from a kd tree into a vector - TODO: expose as iterator instead?
pub fn add_points_to_vector(&self, pts: &mut Vec<VectorN<f64, N>>) {
if let KdTreeImpl::Node(_, ref l, ref pt, ref r) = *self {
l.add_points_to_vector(pts);
pts.push(pt.clone());
r.add_points_to_vector(pts);
}
}
pub fn build(pts: &Vec<VectorN<f64, N>>, dim_idx: usize) -> KdTreeImpl<N>
where
N: DimName,
DefaultAllocator: Allocator<f64, N>,
{
match pts.first() {
Some(first_pt) => {
let next_dim_idx = if dim_idx + 1 >= first_pt.ncols() {
0
} else {
dim_idx + 1
};
match split_vec_with_median(pts, dim_idx) {
(Some(vec1), Some(pt), Some(vec2)) => {
KdTreeImpl::Node(
dim_idx,
Box::new(KdTreeImpl::build(&vec1, next_dim_idx)),
pt,
Box::new(KdTreeImpl::build(&vec2, next_dim_idx)),
)
}
(Some(vec1), Some(pt), None) => {
KdTreeImpl::Node(
dim_idx,
Box::new(KdTreeImpl::build(&vec1, next_dim_idx)),
pt,
Box::new(KdTreeImpl::Empty()),
)
}
(None, Some(pt), Some(vec2)) => {
KdTreeImpl::Node(
dim_idx,
Box::new(KdTreeImpl::Empty()),
pt,
Box::new(KdTreeImpl::build(&vec2, next_dim_idx)),
)
}
(None, Some(pt), None) => {
KdTreeImpl::Node(
dim_idx,
Box::new(KdTreeImpl::Empty()),
pt,
Box::new(KdTreeImpl::Empty()),
)
}
(None, None, None) => {
panic!(
"Should be impossible to have a first point but still split into (None, None, None)."
)
}
(_, None, _) => {
panic!("Meaningless result: point must be contained in a kd-tree node.")
}
}
}
None => KdTreeImpl::Empty(),
}
}
}
fn split_vec_with_median<N>(
pts: &Vec<VectorN<f64, N>>,
dim_idx: usize,
) -> (Option<Vec<VectorN<f64, N>>>, Option<VectorN<f64, N>>, Option<Vec<VectorN<f64, N>>>)
where
N: DimName,
DefaultAllocator: Allocator<f64, N>,
{
unsafe {
let mut pts = pts.to_vec();
pts.sort_by(|a, b| {
a.get_unchecked(0, dim_idx)
.partial_cmp(b.get_unchecked(0, dim_idx))
.unwrap()
});
if pts.len() >= 3 {
let length = pts.len();
let mut middle = pts.split_off(length / 2);
let right = middle.split_off(1);
// (Some(pts), Some(middle[0]), Some(right))
panic!()
} else if pts.len() == 2 {
let middle = pts.split_off(1);
// (Some(pts), Some(middle[0]), None)
panic!()
} else if pts.len() == 1 {
// (None, Some(pts[0]), None)
panic!()
} else {
(None, None, None)
}
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment