kdtree/src/lib.rs

use std::cmp::Ordering;
use std::fmt;

pub enum Side {
    Left,
    Right,
}

/// A k-dimension kd tree
pub enum KdTree<T> {
    Node {
        location: T,
        left: Box<KdTree<T>>,
        right: Box<KdTree<T>>,
    },
    Leaf(Vec<T>),
}

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::min_leaf_size_with_depth(points, dimension, &comparator, leaf_size, 0)
    }

    fn min_leaf_size_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
        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 there are not enough elements for a leaf, make a leaf
        if len <= leaf_size {
            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::min_leaf_size_with_depth(
                &mut left,
                dimension,
                comparator,
                leaf_size,
                depth + 1,
            )),
            right: Box::new(KdTree::min_leaf_size_with_depth(
                right,
                dimension,
                comparator,
                leaf_size,
                depth + 1,
            )),
        }
    }

    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)>,
    ) {
        // Compute the full bounding box of the points
        match *self {
            KdTree::Node {
                ref location,
                ref left,
                ref right,
            } => {
                locations.push((location.clone(), Side::Left));
                left.traverse_leaves_aux(callback, locations);
                locations.pop();

                locations.push((location.clone(), Side::Right));
                right.traverse_leaves_aux(callback, locations);
                locations.pop();
            }
            KdTree::Leaf(ref elements) => {
                callback(locations, elements);
            }
        }
    }
}

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,
            } => {
                writeln!(formatter, "{}{}", indent, location)?;
                left.print(formatter, &format!("{}  ", indent))?;
                right.print(formatter, &format!("{}  ", indent))?;
            }

            &KdTree::Leaf(ref elements) => {
                write!(formatter, "{}[", indent)?;
                for element in elements {
                    write!(formatter, "{}, ", element)?;
                }
                writeln!(formatter, "]")?;
            }
        }

        Ok(())
    }
}

impl<T: fmt::Display> fmt::Display for KdTree<T> {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        self.print(formatter, "")
    }
}