spatial-neighbors 0.1.0

Simple to use implementation of certain Spatial-Partitioning Algorithms | DataStructures
Documentation
use crate::SpatialPartitioner;
use crate::util::in_range;

pub struct QuadTree<Data: Copy> {
    node: QuadTreeNode<Data>,
    capacity: u16,

    x: (f64, f64),
    y: (f64, f64),

    count: usize,
}

impl<Data: Copy> QuadTree<Data> {
    ///
    /// # Arguments
    ///
    /// * `x`: (min_x, max_x) defines the area in wich points can be inserted
    /// * `y`: (min_y, max_y) defines the area in wich points can be inserted
    /// * `capacity`: capacity of each TreeNode
    ///
    pub fn new(x: (f64, f64), y: (f64, f64), capacity: u16) -> QuadTree<Data> {
        QuadTree {
            node: QuadTreeNode::new(((x.1 - x.0) / 2.0, (y.1 - y.0) / 2.0), ((x.1 - x.0), (y.1 - y.0)), capacity),
            x,
            y,
            count: 0,
            capacity,
        }
    }
}

impl<Data: Copy> SpatialPartitioner<Data> for QuadTree<Data> {
    fn insert(&mut self, position: (f64, f64), data: Data) {
        if position.0 < self.x.0 || position.0 >= self.x.1 || position.1 < self.y.0 || position.1 >= self.y.1 {
            panic!("tried to insert position into QuadTree which was out of bounce")
        }

        self.insert_unchecked(position, data);
    }

    fn insert_unchecked(&mut self, position: (f64, f64), data: Data) {
        self.count += 1;
        self.node.insert(position, data);
    }

    fn count(&self) -> usize {
        self.count
    }

    fn clear(&mut self) {
        self.node = QuadTreeNode::new(((self.x.1 - self.x.0) / 2.0, (self.y.1 - self.y.0) / 2.0), ((self.x.1 - self.x.0), (self.y.1 - self.y.0)), self.capacity);
        self.count = 0;
    }

    fn in_circle(&self, position: (f64, f64), radius: f64) -> Vec<Data> {
        let mut data = Vec::new();

        self.node.in_circle(position, radius, &mut data);

        data
    }
}

struct QuadTreeNode<Data: Copy> {
    data: Vec<((f64, f64), Data)>,
    capacity: u16,
    is_full: bool,

    nodes: Option<Box<[QuadTreeNode<Data>; 4]>>,

    center: (f64, f64),
    size: (f64, f64),
}

impl<Data: Copy> QuadTreeNode<Data> {
    fn new(center: (f64, f64), size: (f64, f64), capacity: u16) -> QuadTreeNode<Data> {
        QuadTreeNode {
            data: Vec::new(),
            nodes: None,
            center,
            size,
            capacity,
            is_full: false,
        }
    }

    fn insert(&mut self, location: (f64, f64), data: Data) {
        if !self.is_full {
            self.data.push((location, data));

            if self.data.len() >= self.capacity as usize {
                self.is_full = true;

                self.nodes = Some(Box::new([
                    QuadTreeNode::new((self.center.0 - self.size.0 / 2.0, self.center.1 - self.size.1 / 2.0), (self.size.0 / 2.0, self.size.1 / 2.0), self.capacity),
                    QuadTreeNode::new((self.center.0 - self.size.0 / 2.0, self.center.1 + self.size.1 / 2.0), (self.size.0 / 2.0, self.size.1 / 2.0), self.capacity),
                    QuadTreeNode::new((self.center.0 + self.size.0 / 2.0, self.center.1 - self.size.1 / 2.0), (self.size.0 / 2.0, self.size.1 / 2.0), self.capacity),
                    QuadTreeNode::new((self.center.0 + self.size.0 / 2.0, self.center.1 + self.size.1 / 2.0), (self.size.0 / 2.0, self.size.1 / 2.0), self.capacity),
                ]))
            }
        } else {
            let i = self.get_index(location);
            self.nodes.as_mut().unwrap()[i].insert(location, data);
        }
    }

    fn in_circle(&self, position: (f64, f64), radius: f64, data_vec: &mut Vec<Data>) {
        for data in &self.data {
            if in_range(data.0, position, radius) {
                data_vec.push(data.1)
            }
        }

        if self.nodes.is_none() {
            return;
        }

        let mut indexes = [false; 4];

        indexes[self.get_index((position.0 - radius, position.1 - radius))] = true;
        indexes[self.get_index((position.0 - radius, position.1 + radius))] = true;
        indexes[self.get_index((position.0 + radius, position.1 - radius))] = true;
        indexes[self.get_index((position.0 + radius, position.1 + radius))] = true;

        for (i, index) in indexes.iter().enumerate() {
            if *index {
                self.nodes.as_ref().unwrap()[i].in_circle(position, radius, data_vec);
            }
        }
    }

    fn get_index(&self, location: (f64, f64)) -> usize {
        if self.center.0 > location.0 {
            if self.center.1 > location.1 {
                0
            } else {
                1
            }
        } else if self.center.1 > location.1 {
            2
        } else {
            3
        }
    }
}