Data commit

Task/K-d-tree/Rust/k-d-tree.rust

@@ -0,0 +1,264 @@
+use std::cmp::Ordering;
+use std::cmp::Ordering::Less;
+use std::ops::Sub;
+use std::time::Instant;
+
+use rand::prelude::*;
+
+#[derive(Clone, PartialEq, Debug)]
+struct Point {
+    pub coords: Vec<f32>,
+}
+
+impl<'a, 'b> Sub<&'b Point> for &'a Point {
+    type Output = Point;
+
+    fn sub(self, rhs: &Point) -> Point {
+        assert_eq!(self.coords.len(), rhs.coords.len());
+        Point {
+            coords: self
+                .coords
+                .iter()
+                .zip(rhs.coords.iter())
+                .map(|(&x, &y)| x - y)
+                .collect(),
+        }
+    }
+}
+
+impl Point {
+    fn norm_sq(&self) -> f32 {
+        self.coords.iter().map(|n| n * n).sum()
+    }
+}
+
+struct KDTreeNode {
+    point: Point,
+    dim: usize,
+    // Construction could become faster if we use an arena allocator,
+    // but this is easier to use.
+    left: Option<Box<KDTreeNode>>,
+    right: Option<Box<KDTreeNode>>,
+}
+
+impl KDTreeNode {
+    /// Create a new KDTreeNode around the `dim`th dimension.
+    /// Alternatively, we could dynamically determine the dimension to
+    /// split on by using the longest dimension.
+    pub fn new(points: &mut [Point], dim: usize) -> KDTreeNode {
+        let points_len = points.len();
+        if points_len == 1 {
+            return KDTreeNode {
+                point: points[0].clone(),
+                dim,
+                left: None,
+                right: None,
+            };
+        }
+
+        // Split around the median
+        let pivot = quickselect_by(points, points_len / 2, &|a, b| {
+            a.coords[dim].partial_cmp(&b.coords[dim]).unwrap()
+        });
+
+        let left = Some(Box::new(KDTreeNode::new(
+            &mut points[0..points_len / 2],
+            (dim + 1) % pivot.coords.len(),
+        )));
+        let right = if points.len() >= 3 {
+            Some(Box::new(KDTreeNode::new(
+                &mut points[points_len / 2 + 1..points_len],
+                (dim + 1) % pivot.coords.len(),
+            )))
+        } else {
+            None
+        };
+
+        KDTreeNode {
+            point: pivot,
+            dim,
+            left,
+            right,
+        }
+    }
+
+    pub fn find_nearest_neighbor<'a>(&'a self, point: &Point) -> (&'a Point, usize) {
+        self.find_nearest_neighbor_helper(point, &self.point, (point - &self.point).norm_sq(), 1)
+    }
+
+    fn find_nearest_neighbor_helper<'a>(
+        &'a self,
+        point: &Point,
+        best: &'a Point,
+        best_dist_sq: f32,
+        n_visited: usize,
+    ) -> (&'a Point, usize) {
+        let mut my_best = best;
+        let mut my_best_dist_sq = best_dist_sq;
+        let mut my_n_visited = n_visited;
+
+        // We should always examine the near side
+        if self.point.coords[self.dim] < point.coords[self.dim] && self.right.is_some() {
+            let (a, b) = self.right.as_ref().unwrap().find_nearest_neighbor_helper(
+                point,
+                my_best,
+                my_best_dist_sq,
+                my_n_visited,
+            );
+            my_best = a;
+            my_n_visited = b;
+        } else if self.left.is_some() {
+            let (a, b) = self.left.as_ref().unwrap().find_nearest_neighbor_helper(
+                point,
+                my_best,
+                my_best_dist_sq,
+                my_n_visited,
+            );
+            my_best = a;
+            my_n_visited = b;
+        }
+
+        // distance along this node's axis
+        let axis_dist_sq = (self.point.coords[self.dim] - point.coords[self.dim]).powi(2);
+        if axis_dist_sq <= my_best_dist_sq {
+            // self can only be nearer than best if axis_dist_sq is less than
+            // best_dist_sq because axis_dist_sq is a lower bound for
+            // self_dist_sq
+            let self_dist_sq = (point - &self.point).norm_sq();
+            if self_dist_sq < my_best_dist_sq {
+                my_best = &self.point;
+                my_best_dist_sq = self_dist_sq;
+            }
+
+            // bookkeeping
+            my_n_visited += 1;
+
+            // same reasoning applies for the far side of the split
+            if self.point.coords[self.dim] < point.coords[self.dim] && self.left.is_some() {
+                let (a, b) = self.left.as_ref().unwrap().find_nearest_neighbor_helper(
+                    point,
+                    my_best,
+                    my_best_dist_sq,
+                    my_n_visited,
+                );
+                my_best = a;
+                my_n_visited = b;
+            } else if self.right.is_some() {
+                let (a, b) = self.right.as_ref().unwrap().find_nearest_neighbor_helper(
+                    point,
+                    my_best,
+                    my_best_dist_sq,
+                    my_n_visited,
+                );
+                my_best = a;
+                my_n_visited = b;
+            }
+        }
+
+        (my_best, my_n_visited)
+    }
+}
+
+pub fn main() {
+    let mut rng = thread_rng();
+
+    // wordpress
+    let mut wp_points: Vec<Point> = [
+        [2.0, 3.0],
+        [5.0, 4.0],
+        [9.0, 6.0],
+        [4.0, 7.0],
+        [8.0, 1.0],
+        [7.0, 2.0],
+    ]
+    .iter()
+    .map(|x| Point { coords: x.to_vec() })
+    .collect();
+    let wp_tree = KDTreeNode::new(&mut wp_points, 0);
+
+    let wp_target = Point {
+        coords: vec![9.0, 2.0],
+    };
+    let (point, n_visited) = wp_tree.find_nearest_neighbor(&wp_target);
+    println!("Wikipedia example data:");
+    println!("Point: [9, 2]");
+    println!("Nearest neighbor: {:?}", point);
+    println!("Distance: {}", (point - &wp_target).norm_sq().sqrt());
+    println!("Nodes visited: {}", n_visited);
+
+    // randomly generated 3D
+    let n_random = 1000;
+    let mut make_random_point = || Point {
+        coords: (0..3).map(|_| (rng.gen::<f32>() - 0.5) * 1000.0).collect(),
+    };
+    let mut random_points: Vec<Point> = (0..n_random).map(|_| make_random_point()).collect();
+
+    let start_cons_time = Instant::now();
+    let random_tree = KDTreeNode::new(&mut random_points, 0);
+    let cons_time = start_cons_time.elapsed();
+    println!(
+        "1,000 3d points (Construction time: {}ms)",
+        cons_time.as_millis(),
+    );
+
+    let random_target = make_random_point();
+
+    let (point, n_visited) = random_tree.find_nearest_neighbor(&random_target);
+    println!("Point: {:?}", random_target);
+    println!("Nearest neighbor: {:?}", point);
+    println!("Distance: {}", (point - &random_target).norm_sq().sqrt());
+    println!("Nodes visited: {}", n_visited);
+
+    // benchmark search time
+    let n_searches = 1000;
+    let random_targets: Vec<Point> = (0..n_searches).map(|_| make_random_point()).collect();
+
+    let start_search_time = Instant::now();
+    let mut total_n_visited = 0;
+    for target in &random_targets {
+        let (_, n_visited) = random_tree.find_nearest_neighbor(target);
+        total_n_visited += n_visited;
+    }
+    let search_time = start_search_time.elapsed();
+    println!(
+        "Visited an average of {} nodes on {} searches in {} ms",
+        total_n_visited as f32 / n_searches as f32,
+        n_searches,
+        search_time.as_millis(),
+    );
+}
+
+fn quickselect_by<T>(arr: &mut [T], position: usize, cmp: &dyn Fn(&T, &T) -> Ordering) -> T
+where
+    T: Clone,
+{
+    // We use `thread_rng` here because it was already initialized in `main`.
+    let mut pivot_index = thread_rng().gen_range(0, arr.len());
+    // Need to wrap in another closure or we get ownership complaints.
+    // Tried using an unboxed closure to get around this but couldn't get it to work.
+    pivot_index = partition_by(arr, pivot_index, &|a: &T, b: &T| cmp(a, b));
+    let array_len = arr.len();
+    match position.cmp(&pivot_index) {
+        Ordering::Equal => arr[position].clone(),
+        Ordering::Less => quickselect_by(&mut arr[0..pivot_index], position, cmp),
+        Ordering::Greater => quickselect_by(
+            &mut arr[pivot_index + 1..array_len],
+            position - pivot_index - 1,
+            cmp,
+        ),
+    }
+}
+
+fn partition_by<T>(arr: &mut [T], pivot_index: usize, cmp: &dyn Fn(&T, &T) -> Ordering) -> usize {
+    let array_len = arr.len();
+    arr.swap(pivot_index, array_len - 1);
+    let mut store_index = 0;
+    for i in 0..array_len - 1 {
+        if cmp(&arr[i], &arr[array_len - 1]) == Less {
+            arr.swap(i, store_index);
+            store_index += 1;
+        }
+    }
+    arr.swap(array_len - 1, store_index);
+    store_index
+}