I have a simple QuadTree. For readability and to help me understand what is happening, It avoids recursion and has a static depth. The QuadTree stores references to the points that are owned by another container.
struct QuadLeaf<'a> {
vec: Vec<&'a (f32,f32)>,
rect: (f32, f32, f32, f32)
}
struct QuadTwig<'a> {
cells: [QuadLeaf<'a>; 4],
}
struct QuadBranch<'a> {
cells: [QuadTwig<'a>; 4],
}
struct QuadTree<'a> {
cells: [QuadBranch<'a>; 4],
}
Constructing and inserting into this tree is relatively simple. The QuadLeaf is constructed with a bounding rect and an empty vec, and has a method that attempts to insert a point. It returns true if the point is within the rect and has been inserted.
impl<'a> QuadLeaf<'a> {
fn new(rect: (f32, f32, f32, f32)) -> Self {
QuadLeaf {vec: Vec::new(),rect}
}
fn insert(&mut self, point: &'a (f32, f32)) -> bool {
if is_point_in_rect(point.0, point.1, self.rect) {
self.vec.push(point);
true
} else {
false
}
}
}
The QuadTwig new function splits the bounding rect into 4 and creates 4 new QuadLeafs . It's insert function attempts to insert into each leaf, short circuiting when it is successful, and returning false if unsuccessful.
impl<'a> QuadTwig<'a> {
fn new(rect: (f32, f32, f32, f32)) -> Self {
let rects = divide_into_4(rect);
QuadTwig {
cells: [
QuadLeaf::new(rects[0]),
QuadLeaf::new(rects[1]),
QuadLeaf::new(rects[2]),
QuadLeaf::new(rects[3])
]
}
}
fn insert(&mut self, point: &'a (f32, f32)) -> bool {
for cell in self.cells.iter_mut() {
if cell.insert(point) {
return true;
}
}
false
}
}
The implementations for QuadBranch and QuadTree are exactly the same, but the new function just constructs the next level down in the tree. This could be refactored later for less code duplication, but for demonstration purposes I will leave it. I also think it does not matter for the context of this question.
Question:
I want to create an Iterator that yields each point in the tree, and the 9 leaves that it is close to (or inside of).
I have managed to create a simpler version, that just yields each point and the leaf it is in:
/// An Iterator that yields each point and the leaf it is in
struct PointAndLeafIterator<'a> {
ptr: &'a QuadTree<'a>,
index: (usize, usize, usize, usize)
}
/// An Iterator that yields each point and the leaf it is in
impl<'a> Iterator for PointAndLeafIterator<'a> {
/// Returns (point, leaf)
type Item = (&'a (f32, f32), Vec<&'a (f32, f32)>);
/// Starts at (0,0,0,0) and ends at (3, 3, 3, num_points_in_leaf)
/// It increases the index by 1 each time, and if it reaches the end of the cell, it moves to the next cell
fn next(&mut self) -> Option<Self::Item> {
let (branch_index, twig_index, leaf_index, point_index) = &mut self.index;
let branch = &self.ptr.cells[*branch_index];
let twig = &branch.cells[*twig_index];
let leaf = &twig.cells[*leaf_index];
let point = leaf.vec.get(*point_index);
//go through all the points in the leaf
if let Some(point) = point {
*point_index += 1;
return Some((point, leaf.vec.clone()));
}
//if we reach the end of the leaf, go to the next leaf
*point_index = 0;
*leaf_index += 1;
if *leaf_index < 4 {
return self.next();
}
//if we reach the end of the twig, go to the next twig
*leaf_index = 0;
*twig_index += 1;
if *twig_index < 4 {
return self.next();
}
//if we reach the end of the branch, go to the next branch
*twig_index = 0;
*branch_index += 1;
if *branch_index < 4 {
return self.next();
}
//if we reach the end of the tree, we are done
None
}
}
This can be used like this:
fn main() {
let points: Vec<(f32, f32)> = vec![
(0.0, 0.0),
(1.0, 1.0),
(31.0,31.0),
(2.0, 2.0),
(3.0, 3.0),
(32.0,32.0),
];
let mut quadtree = QuadTree::new((0.0, 0.0, 40.0, 40.0));
for point in points.iter() {
quadtree.insert(point);
}
for (point, leaf) in quadtree.into_point_and_leaf_iter() {
println!("Point: {:?}", point);
println!("Leaf: {:?}", leaf);
}
}
However the neighbouring version is proving to be much more difficult. How can I write this algorithm?
/// An Iterator that yields each point, the leaf it is in, and the neighbouring leaves
struct PointAndLeafAndNeighboursIterator<'a> {
ptr: &'a QuadTree<'a>,
index: (usize, usize, usize, usize)
}
impl<'a> Iterator for PointAndLeafAndNeighboursIterator<'a> {
///Return the 9 leaves that surround the point
///If there is no leaf in a direction, it will return an empty leaf
type Item = (&'a (f32, f32), [Vec<&'a (f32, f32)>; 9]);
/// Starts at (0,0,0,0) and ends at (3, 3, 3, num_points_in_leaf)
/// It increases the index by 1 each time, and if it reaches the end of the cell, it moves to the next cell
fn next(&mut self) -> Option<Self::Item> {
unimplemented!()
}
}
Here is a Rust playground link to all code in this question.
First I added indexing by leaf indices since the QuadTree struct is basically an 8x8 matrix of leaves.
use std::ops::Index;
#[derive(Clone, Copy, Debug)]
struct QuadLeafId {
x: i8,
y: i8,
}
impl QuadLeafId {
fn new(x: i8, y: i8) -> Self {
Self { x, y }
}
fn level_index(&self, level: usize) -> usize {
(((self.y >> level) % 2) * 2 + ((self.x >> level) % 2)) as usize
}
/// Extract the QuadTwig array index containing this leaf
fn twig_index(&self) -> usize {
self.level_index(0)
}
/// Extract the QuadBranch array index containing this leaf
fn branch_index(&self) -> usize {
self.level_index(1)
}
/// Extract the QuadTree array index containing this leaf
fn tree_index(&self) -> usize {
self.level_index(2)
}
}
impl<'a> Index<QuadLeafId> for QuadTwig<'a> {
type Output = QuadLeaf<'a>;
fn index(&self, index: QuadLeafId) -> &Self::Output {
&self.cells[index.twig_index()]
}
}
impl<'a> Index<QuadLeafId> for QuadBranch<'a> {
type Output = QuadLeaf<'a>;
fn index(&self, index: QuadLeafId) -> &Self::Output {
&self.cells[index.branch_index()][index]
}
}
impl<'a> Index<QuadLeafId> for QuadTree<'a> {
type Output = QuadLeaf<'a>;
fn index(&self, index: QuadLeafId) -> &Self::Output {
&self.cells[index.tree_index()][index]
}
}
Then added iterating over nearby leaves in a QuadTree
impl QuadLeafId {
/// The ids of the 9 nearby leaves
fn near_ids(self) -> impl Iterator<Item = Self> {
(-1..=1).flat_map(move |x| {
(-1..=1).map(move |y| Self {
x: self.x + x,
y: self.y + y,
})
})
}
fn is_valid(&self) -> bool {
0 <= self.y && self.y < 8 && 0 <= self.x && self.x < 8
}
}
impl<'a> QuadTree<'a> {
fn get_leaf(&self, id: QuadLeafId) -> Option<&QuadLeaf<'a>> {
id.is_valid().then(|| &self[id])
}
fn near_leaves(&self, id: QuadLeafId) -> impl Iterator<Item = Option<&QuadLeaf<'a>>> {
id.near_ids().map(|id| self.get_leaf(id))
}
}
Implementing the iterator afterwards is straightforward:
impl QuadLeafId {
fn next_id(mut self) -> Option<Self> {
self.x += 1;
if self.x < 8 {
return Some(self);
}
self.x = 0;
self.y += 1;
(self.y < 8).then_some(self)
}
}
impl<'a> QuadTree<'a> {
fn into_point_and_leaf_and_neighbours_iter(
&'a mut self,
) -> PointAndLeafAndNeighboursIterator<'a> {
PointAndLeafAndNeighboursIterator::<'a> {
ptr: self,
index: Some(QuadLeafId::new(0, 0)),
point_index: 0,
}
}
}
/// An Iterator that yields each point, the leaf it is in, and the neighboring leaves
struct PointAndLeafAndNeighboursIterator<'a> {
ptr: &'a QuadTree<'a>,
index: Option<QuadLeafId>,
point_index: usize,
}
impl<'a> Iterator for PointAndLeafAndNeighboursIterator<'a> {
///Return the 9 leaves that surround the point
///If there is no leaf in a direction, it will return an empty leaf
type Item = (&'a (f32, f32), [Vec<&'a (f32, f32)>; 9]);
/// Starts at (0,0,0,0) and ends at (3, 3, 3, num_points_in_leaf)
/// It increases the index by 1 each time, and if it reaches the end of the cell, it moves to the next cell
fn next(&mut self) -> Option<Self::Item> {
loop {
let ind = self.index?;
let leaf = &self.ptr[ind];
if let Some(point) = leaf.vec.get(self.point_index) {
self.point_index += 1;
let mut iter = self
.ptr
.near_leaves(ind)
.map(|leaf| leaf.map(|leaf| leaf.vec.clone()).unwrap_or_default());
return Some((*point, [(); 9].map(|_| iter.next().unwrap())));
}
self.point_index = 0;
self.index = ind.next_id();
}
}
}