fux_kdtree 0.2.0

K-dimensional tree implemented in Rust for fast NN querying.
Documentation
use ::kdtree::*;

enum PointsWereOnSide {
    Left,
    Right,
    Both,
}

struct PartitionPointHelper {
    points_were_on_side: PointsWereOnSide,
    index_of_splitter: usize,
}

fn partition_sliding_midpoint_helper<T: KdtreePointTrait>(vec: &mut [T], midpoint_value: f64, partition_on_dimension: usize) -> PartitionPointHelper {
    let mut closest_index = 0;
    let mut closest_distance = (vec[0].dims()[partition_on_dimension] - midpoint_value).abs();

    const HAS_POINTS_ON_LEFT_SIDE: i32 = 0b01;
    const HAS_POINTS_ON_RIGHT_SIDE: i32 = 0b10;
    const HAS_POINTS_ON_BOTH_SIDES: i32 = HAS_POINTS_ON_RIGHT_SIDE | HAS_POINTS_ON_LEFT_SIDE;
    let mut has_points_on_sides = 0;

    for i in 0..vec.len() {
        let p = vec.get(i).unwrap();
        if p.dims()[partition_on_dimension] <= midpoint_value {
            has_points_on_sides |= HAS_POINTS_ON_LEFT_SIDE;
        } else {
            has_points_on_sides |= HAS_POINTS_ON_RIGHT_SIDE;
        }

        let dist = (p.dims()[partition_on_dimension] - midpoint_value).abs();

        if dist < closest_distance {
            closest_distance = dist;
            closest_index = i;
        }
    }

    if has_points_on_sides != HAS_POINTS_ON_BOTH_SIDES {
        return PartitionPointHelper {
            index_of_splitter: closest_index,
            points_were_on_side: if has_points_on_sides == HAS_POINTS_ON_LEFT_SIDE { PointsWereOnSide::Left } else { PointsWereOnSide::Right }
        }
    }
    return PartitionPointHelper {
        index_of_splitter: closest_index,
        points_were_on_side: PointsWereOnSide::Both
    }
}

pub fn partition_sliding_midpoint<T: KdtreePointTrait>(vec: &mut [T], midpoint_value: f64, partition_on_dimension: usize) -> usize {
    let vec_len = vec.len();
    debug_assert!(vec[0].dims().len() > partition_on_dimension);

    if vec.len() == 1 {
        return 0;
    }

    let partition_point_data = partition_sliding_midpoint_helper(vec, midpoint_value, partition_on_dimension);

    match partition_point_data.points_were_on_side {
        PointsWereOnSide::Left => {
            vec.swap(partition_point_data.index_of_splitter, vec_len - 1);
            vec_len - 1
        },
        PointsWereOnSide::Right => {
            vec.swap(partition_point_data.index_of_splitter, 0);
            0
        }
        PointsWereOnSide::Both => {
            let index_of_splitting_point = partition_kdtree(vec, partition_point_data.index_of_splitter, partition_on_dimension);
            index_of_splitting_point
        }
    }
}

fn partition_kdtree<T: KdtreePointTrait>(vec: &mut [T], index_of_splitting_point: usize, partition_on_dimension: usize) -> usize {
    if vec.len() == 1 {
        return 0;
    }

    let pivot = vec[index_of_splitting_point].dims()[partition_on_dimension];
    let vec_len = vec.len();

    vec.swap(index_of_splitting_point, vec_len - 1);

    let mut left = 0usize;
    let mut right = vec.len() - 2;
    let mut last_succesful_swap = vec.len() - 1;

    //variant of Lomuto algo.
    loop {
        while left <= right && vec[left].dims()[partition_on_dimension] <= pivot {
            left += 1;
        }

        while right > left && vec[right].dims()[partition_on_dimension] > pivot {
            right -= 1;
        }

        if right > left {
            vec.swap(left, right);
            last_succesful_swap = right;

            left += 1;
            right -= 1;
        } else {
            break;
        }
    }

    if last_succesful_swap == vec_len - 1 && vec[right].dims()[partition_on_dimension] > pivot {
        vec.swap(right, last_succesful_swap);
        last_succesful_swap = right;
    } else if vec[left].dims()[partition_on_dimension] > pivot {
        vec.swap(left, vec_len - 1);
        last_succesful_swap = left;
    } else {
        vec.swap(last_succesful_swap, vec_len - 1);
    }

    last_succesful_swap
}


#[cfg(test)]
mod tests {
    use ::kdtree::*;
    use ::kdtree::test_common::*;

    use ::rand::distributions::{IndependentSample, Range};
    use ::rand::*;

    use super::*;
    use super::partition_kdtree;

    #[test]
    fn parition_kdtree_works() {
        let p1 = Point2WithId::new(0, 1., 4.);
        let p2 = Point2WithId::new(1, 2., 6.);
        let p3 = Point2WithId::new(2, 3., 8.);
        let p4 = Point2WithId::new(3, 0., 8.);
        let p5 = Point2WithId::new(4, -1., 8.);
        let p6 = Point2WithId::new(5, 3., 8.);
        let p7 = Point2WithId::new(6, 4., 8.);

        let vec = vec![p1, p2, p3, p4, p5, p6, p7];
        assert_eq! (1, partition_kdtree(&mut vec.clone(), 3, 0));

        assert_eq! (6, partition_kdtree(&mut vec.clone(), 6, 0));

        assert_eq! (0, partition_kdtree(&mut vec.clone(), 4, 0));

        assert_eq! (5, partition_kdtree(&mut vec.clone(), 2, 0));
    }

    quickcheck! {
        fn partition_kdtree_qc(xs: Vec<f64>) -> bool {
            let mut vec : Vec<Point1WithId> = vec![];

            for i in 0 .. xs.len() {
                let p = Point1WithId::new(i as i32, xs[i]);
                vec.push(p);
            }

            if xs.len() == 0 {
                return true;
            }
            let between = Range::new(0, xs.len());
            let mut rng = thread_rng();

            for _ in 0 .. 5 {
                let random_splitting_index = between.ind_sample(&mut rng);

                let mut vec = vec.clone();

                let index_of_splitting_point = partition_kdtree(&mut vec, random_splitting_index, 0);
                return assert_partition(&vec, index_of_splitting_point);
            }

            true
        }
    }

    #[test]
    fn partition_given_midpoint_exactly_in_between_points_returns_smaller_index() {
        let p1 = Point2WithId::new(1, 2., 4.);
        let p2 = Point2WithId::new(1, 4., 6.);
        let mut vec = vec![p1, p2];

        assert_eq! (0, partition_sliding_midpoint(&mut vec, 3., 0));
        assert_eq! (0, partition_sliding_midpoint(&mut vec, 5., 1));
    }

    #[test]
    fn partition_given_midpoint_which_has_all_points_on_one_side_slides_split_plane_and_returns_index_to_closest_element() {
        let p1 = Point2WithId::new(1, 2., 4.);
        let p2 = Point2WithId::new(2, 4., 6.);
        let p3 = Point2WithId::new(3, 3., 7.);
        let p4 = Point2WithId::new(4, 0., 8.);
        let mut vec = vec![p1, p2, p3];

        assert_eq! (0, partition_sliding_midpoint(&mut vec, 1.9, 0));
        assert_eq! (1, vec[0].id);
        assert_eq! (2, vec[1].id);
        assert_eq! (3, vec[2].id);

        let mut vec = vec![p1, p2, p3];
        assert_eq! (0, partition_sliding_midpoint(&mut vec, -5000., 0));
        assert_eq! (1, vec[0].id);
        assert_eq! (2, vec[1].id);
        assert_eq! (3, vec[2].id);

        let mut vec = vec![p1, p2, p3];
        assert_eq! (2, partition_sliding_midpoint(&mut vec, 10., 0));
        assert_eq! (1, vec[0].id);
        assert_eq! (3, vec[1].id);
        assert_eq! (2, vec[2].id);


        let mut vec = vec![p1, p2, p3];
        assert_eq! (2, partition_sliding_midpoint(&mut vec, 10., 1));
        assert_eq! (1, vec[0].id);
        assert_eq! (2, vec[1].id);
        assert_eq! (3, vec[2].id);

        let mut vec = vec![p1, p2, p3, p4];
        assert_eq! (0, partition_sliding_midpoint(&mut vec, -5000., 0));
        assert_eq! (4, vec[0].id);
        assert_eq! (2, vec[1].id);
        assert_eq! (3, vec[2].id);
        assert_eq! (1, vec[3].id);
    }

    fn assert_partition(v: &Vec<Point1WithId>, index_of_splitting_point: usize) -> bool {
        let pivot = v[index_of_splitting_point].dims()[0];

        for i in 0..index_of_splitting_point {
            if v[i].dims()[0] > pivot {
                return false;
            }
        }

        for i in index_of_splitting_point + 1..v.len() {
            if v[i].dims()[0] < pivot {
                return false;
            }
        }

        true
    }
}