我有一个简单的QuadTree。为了可读性和帮助我理解正在发生的事情，它避免了递归并具有静态深度。QuadTree存储对另一个容器拥有的点的引用。
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],
}
构造和插入这个树相对简单。QuadLeaf是由一个有边界的rect和一个空的vec构造的，并且有一个尝试插入一个点的方法。如果该点在rect中并且已经插入，则返回true。
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
}
}
}
QuadTwignew函数将绑定的rect拆分为4个，并创建4个新的QuadLeaf。它的insert函数尝试插入到每个叶子中，成功时短路，不成功则返回false。
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
}
}
QuadBranch和QuadTree的实现是完全相同的，但是new函数只是构造树中的下一层。为了减少代码重复，可以稍后对其进行重构，但出于演示目的，我将保留它。我也认为这与这个问题的背景无关。
问题:
我想创建一个迭代器，生成树中的每个点，以及它接近(或内部)的9个叶子。
我已经设法创建了一个更简单的版本，它只产生每个点和它所在的叶子:
/// 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
}
}
可以这样使用:
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);
}
}
然而，相邻的版本证明要困难得多。我怎么写这个算法呢?
/// 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!()
}
}
下面是这个问题中所有代码的Rust playground链接。
我的实现。
首先，我添加了叶索引索引，因为QuadTree结构体基本上是一个8x8的叶矩阵。
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]
}
}
然后在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))
}
}
之后实现迭代器很简单:
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();
}
}
}