Skip to content

Instantly share code, notes, and snippets.

@djmcgill
Last active October 7, 2017 23:23
Show Gist options
  • Save djmcgill/219c8f10045da8d2dd8e94fd668f0258 to your computer and use it in GitHub Desktop.
Save djmcgill/219c8f10045da8d2dd8e94fd668f0258 to your computer and use it in GitHub Desktop.
kd_tree agaaaiin
use na::*;
use na::allocator::Allocator;
use std::vec::IntoIter;
/// An n-dimensional point with a defined n-vector from the origin
pub trait NPoint: Clone
where
DefaultAllocator: Allocator<f64, Self::N>,
{
type N: DimName;
fn from_origin(&self) -> &VectorN<f64, Self::N>;
}
impl<D: DimName> NPoint for VectorN<f64, D>
where
DefaultAllocator: Allocator<f64, D>,
{
type N = D;
fn from_origin(&self) -> &VectorN<f64, D> {
self
}
}
/// A three dimensional plane defined by a point3 and a vector3
pub struct Plane3 {
position: Vector3<f64>,
normal: Vector3<f64>,
}
/// An arbitrary structure of points capable of answering spatial queries
pub trait SpatialQueryStructure<T: NPoint<N = U3>> {
fn find_closest_point(&self, p: Vector3<f64>) -> T;
fn find_closest_point_within_range(&self, p: Vector3<f64>, range: f64) -> T;
type KNearestIterator: Iterator<Item = T>;
fn find_k_nearest_points(&self, p: Vector3<f64>, k: u32) -> Self::KNearestIterator; // impl Iterator would be better
type KNearestRangeIter: Iterator<Item = T>;
fn find_k_nearest_points_within_range(
&self,
p: Vector3<f64>,
k: u32,
range: f64,
) -> Self::KNearestRangeIter;
}
struct Foo {}
impl SpatialQueryStructure<Vector3<f64>> for Foo {
fn find_closest_point(&self, p: Vector3<f64>) -> Vector3<f64> {
panic!();
}
fn find_closest_point_within_range(&self, p: Vector3<f64>, range: f64) -> Vector3<f64> {
panic!();
}
type KNearestIterator = IntoIter<Vector3<f64>>;
fn find_k_nearest_points(&self, p: Vector3<f64>, k: u32) -> Self::KNearestIterator {
let vec: Vec<Vector3<f64>> = panic!();
vec.into_iter()
}
type KNearestRangeIter = IntoIter<Vector3<f64>>;
fn find_k_nearest_points_within_range(
&self,
p: Vector3<f64>,
k: u32,
range: f64,
) -> Self::KNearestRangeIter {
panic!();
}
}
use na::*;
use na::allocator::Allocator;
use domain::*;
/// A kd tree
pub enum KdTree<P>
where
P: NPoint,
DefaultAllocator: Allocator<f64, P::N>,
{
Node(usize, Box<KdTree<P>>, P, Box<KdTree<P>>),
Empty(),
}
// /// 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
// }
impl<P> KdTree<P>
where
P: NPoint,
DefaultAllocator: Allocator<f64, P::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<P>) {
if let KdTree::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<P>, dim_idx: usize) -> KdTree<P> {
if pts.is_empty() {
return KdTree::Empty();
}
// The dimensionality is statically known
let next_dim_idx = if dim_idx + 1 > P::N::dim() {
0
} else {
dim_idx + 1
};
let (left, point, right) = KdTree::split_vec_with_median(pts, dim_idx).unwrap(); // panics on empty vec
KdTree::Node(
dim_idx,
Box::new(KdTree::build(left, next_dim_idx)),
point,
Box::new(KdTree::build(right, next_dim_idx)),
)
}
/// Returns None on an empty vec. Returns empty l/r subtrees if too few elements
fn split_vec_with_median(mut pts: Vec<P>, dim_idx: usize) -> Option<(Vec<P>, P, Vec<P>)>
where
DefaultAllocator: Allocator<f64, P::N>,
{
pts.sort_by(|a, b| unsafe {
a.from_origin()
.get_unchecked(0, dim_idx)
.partial_cmp(b.from_origin().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);
// `remove` takes the item out of the vec, which is preferable to copying
Some((pts, middle.remove(0), right))
} else if pts.len() == 2 {
let mut middle = pts.split_off(1);
Some((pts, middle.remove(0), Vec::new()))
} else if pts.len() == 1 {
Some((Vec::new(), pts.remove(0), Vec::new()))
} else {
None
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment