Cleaner
This commit is contained in:
parent
64f25f51bd
commit
a4176bca5a
|
@ -4,7 +4,6 @@ version = "0.1.0"
|
||||||
authors = ["Thomas Forgione <thomas@tforgione.fr>"]
|
authors = ["Thomas Forgione <thomas@tforgione.fr>"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
num = "*"
|
|
||||||
|
|
||||||
[[bin]]
|
[[bin]]
|
||||||
name = "kdtree-example"
|
name = "kdtree-example"
|
||||||
|
|
|
@ -1,45 +1,26 @@
|
||||||
extern crate num;
|
|
||||||
extern crate kdtree;
|
extern crate kdtree;
|
||||||
|
|
||||||
use std::fmt;
|
use std::fmt;
|
||||||
use std::ops::Index;
|
|
||||||
|
|
||||||
use kdtree::Zero;
|
|
||||||
use kdtree::KdTree;
|
use kdtree::KdTree;
|
||||||
|
|
||||||
#[derive(Copy, Clone)]
|
#[derive(Clone)]
|
||||||
struct Vector2 {
|
struct Vector2 {
|
||||||
data: [f64; 2],
|
data: [f64; 2],
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Vector2 {
|
impl Vector2 {
|
||||||
pub fn new(x: f64, y: f64) -> Vector2 {
|
pub fn new(x: f64, y: f64) -> Vector2 {
|
||||||
Vector2 {
|
Vector2 { data: [x, y] }
|
||||||
data: [x, y],
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Index<usize> for Vector2 {
|
|
||||||
type Output = f64;
|
|
||||||
fn index(&self, index: usize) -> &f64 {
|
|
||||||
&self.data[index]
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Zero for Vector2 {
|
|
||||||
fn zero() -> Vector2 {
|
|
||||||
Vector2::new(0.0, 0.0)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl fmt::Display for Vector2 {
|
impl fmt::Display for Vector2 {
|
||||||
fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
|
fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
|
||||||
write!(formatter, "({}, {})", self[0], self[1])
|
write!(formatter, "({}, {})", self.data[0], self.data[1])
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
let mut points = vec![
|
let mut points = vec![
|
||||||
Vector2::new(2.0, 3.0),
|
Vector2::new(2.0, 3.0),
|
||||||
|
@ -50,8 +31,12 @@ fn main() {
|
||||||
Vector2::new(7.0, 2.0),
|
Vector2::new(7.0, 2.0),
|
||||||
];
|
];
|
||||||
|
|
||||||
let kdtree = KdTree::new(&mut points, 2, 1);
|
let kdtree = KdTree::new(
|
||||||
|
&mut points,
|
||||||
|
2,
|
||||||
|
|coord, ref a, ref b| a.data[coord].partial_cmp(&b.data[coord]).unwrap(),
|
||||||
|
1,
|
||||||
|
);
|
||||||
|
|
||||||
println!("{}", kdtree);
|
println!("{}", kdtree);
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
91
src/lib.rs
91
src/lib.rs
|
@ -1,76 +1,96 @@
|
||||||
extern crate num;
|
|
||||||
|
|
||||||
use std::fmt;
|
use std::fmt;
|
||||||
use std::ops::Index;
|
use std::cmp::Ordering;
|
||||||
use std::cmp::PartialOrd;
|
|
||||||
|
|
||||||
pub trait Zero {
|
|
||||||
fn zero() -> Self;
|
|
||||||
}
|
|
||||||
|
|
||||||
impl<T: num::Zero> Zero for T {
|
|
||||||
fn zero() -> Self {
|
|
||||||
num::Zero::zero()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// A k-dimension kd tree
|
/// A k-dimension kd tree
|
||||||
pub enum KdTree<T> {
|
pub enum KdTree<T> {
|
||||||
Node { location: T, left: Box<KdTree<T>>, right: Box<KdTree<T>> },
|
Node {
|
||||||
|
location: T,
|
||||||
|
left: Box<KdTree<T>>,
|
||||||
|
right: Box<KdTree<T>>,
|
||||||
|
},
|
||||||
Leaf(Vec<T>),
|
Leaf(Vec<T>),
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<T, U> KdTree<T> where T: Copy + Clone + Index<usize, Output = U> + Zero, U: PartialOrd {
|
impl<T: Clone> KdTree<T> where {
|
||||||
pub fn new(points: &mut [T], dimension: usize, leaf_size: usize) -> KdTree<T> {
|
pub fn new<F: Fn(usize, &T, &T) -> Ordering>(
|
||||||
KdTree::new_with_depth(points, dimension, leaf_size, 0)
|
points: &mut [T],
|
||||||
|
dimension: usize,
|
||||||
|
comparator: F,
|
||||||
|
leaf_size: usize,
|
||||||
|
) -> KdTree<T> {
|
||||||
|
KdTree::new_with_depth(points, dimension, &comparator, leaf_size, 0)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn new_with_depth(points: &mut [T], dimension: usize, leaf_size: usize, depth: usize) -> KdTree<T> {
|
fn new_with_depth<F: Fn(usize, &T, &T) -> Ordering>(
|
||||||
|
points: &mut [T],
|
||||||
|
dimension: usize,
|
||||||
|
comparator: &F,
|
||||||
|
leaf_size: usize,
|
||||||
|
depth: usize,
|
||||||
|
) -> KdTree<T> {
|
||||||
// Select axis to work on
|
// Select axis to work on
|
||||||
let axis = depth % dimension;
|
let axis = depth % dimension;
|
||||||
|
|
||||||
// Sort the points by axis
|
// Sort the points by axis
|
||||||
points.sort_by(|p1, p2| p1[axis].partial_cmp(&p2[axis]).unwrap());
|
points.sort_by(|a, b| comparator(axis, a, b));
|
||||||
|
|
||||||
// Find the median
|
// Find the median
|
||||||
let len = points.len();
|
let len = points.len();
|
||||||
|
|
||||||
// If there are not enough elements for a leaf, make a leaf
|
// If there are not enough elements for a leaf, make a leaf
|
||||||
if len <= leaf_size {
|
if len <= leaf_size {
|
||||||
let mut contained = vec![Zero::zero(); len];
|
let mut contained = vec![];
|
||||||
contained.copy_from_slice(points);
|
for p in points {
|
||||||
|
contained.push(p.clone());
|
||||||
|
}
|
||||||
return KdTree::Leaf(contained);
|
return KdTree::Leaf(contained);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Else split what's remaining
|
// Else split what's remaining
|
||||||
let delimiter = points.len() / 2;
|
let delimiter = points.len() / 2;
|
||||||
let median = points[delimiter];
|
let median = points[delimiter].clone();
|
||||||
|
|
||||||
// Build the left and right branches
|
// Build the left and right branches
|
||||||
let zero = Zero::zero();
|
let mut left = vec![];
|
||||||
let mut left = vec![zero; delimiter];
|
for i in &points[0..delimiter] {
|
||||||
left.copy_from_slice(&points[0..delimiter]);
|
left.push(i.clone());
|
||||||
|
}
|
||||||
|
|
||||||
let right = &mut points[delimiter + 1 .. len];
|
let right = &mut points[delimiter + 1..len];
|
||||||
|
|
||||||
// Recursive call
|
// Recursive call
|
||||||
KdTree::Node {
|
KdTree::Node {
|
||||||
location: median,
|
location: median,
|
||||||
left: Box::new(KdTree::new_with_depth(&mut left, dimension, leaf_size, depth + 1)),
|
left: Box::new(KdTree::new_with_depth(
|
||||||
right: Box::new(KdTree::new_with_depth(right, dimension, leaf_size, depth + 1)),
|
&mut left,
|
||||||
|
dimension,
|
||||||
|
comparator,
|
||||||
|
leaf_size,
|
||||||
|
depth + 1,
|
||||||
|
)),
|
||||||
|
right: Box::new(KdTree::new_with_depth(
|
||||||
|
right,
|
||||||
|
dimension,
|
||||||
|
comparator,
|
||||||
|
leaf_size,
|
||||||
|
depth + 1,
|
||||||
|
)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<T: Zero + Copy + Clone + fmt::Display + Index<usize>> 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 { 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))?;
|
||||||
},
|
}
|
||||||
|
|
||||||
&KdTree::Leaf(ref elements) => {
|
&KdTree::Leaf(ref elements) => {
|
||||||
write!(formatter, "{}[", indent)?;
|
write!(formatter, "{}[", indent)?;
|
||||||
|
@ -78,16 +98,15 @@ impl<T: Zero + Copy + Clone + fmt::Display + Index<usize>> KdTree<T> {
|
||||||
write!(formatter, "{}, ", element)?;
|
write!(formatter, "{}, ", element)?;
|
||||||
}
|
}
|
||||||
writeln!(formatter, "]")?;
|
writeln!(formatter, "]")?;
|
||||||
},
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<T: Zero + Copy + Clone + fmt::Display + Index<usize>> fmt::Display for KdTree<T> {
|
impl<T: fmt::Display> fmt::Display for KdTree<T> {
|
||||||
fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
|
fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
|
||||||
self.print(formatter, "")
|
self.print(formatter, "")
|
||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue