如何为静态深度非递归QuadTree实现含相邻叶子的迭代器

问题背景 #

已实现静态深度、非递归的QuadTree，存储外部容器中点的引用，完成了PointAndLeafIterator（返回每个点及其所在叶子）。现在需要实现PointAndLeafAndNeighboursIterator，使其返回每个点、所在叶子，以及周围9个相邻叶子（无对应叶子时返回空Vec）。

完整解决方案 #

1. 补充必要的辅助函数与层级实现 #

先补充QuadBranch、QuadTree的实现，以及坐标判断、矩形分割的辅助函数：

fn is_point_in_rect(x: f32, y: f32, rect: (f32, f32, f32, f32)) -> bool {
    let (x_min, y_min, x_max, y_max) = rect;
    x >= x_min && x <= x_max && y >= y_min && y <= y_max
}

fn divide_into_4(rect: (f32, f32, f32, f32)) -> [(f32, f32, f32, f32); 4] {
    let (x_min, y_min, x_max, y_max) = rect;
    let x_mid = (x_min + x_max) / 2.0;
    let y_mid = (y_min + y_max) / 2.0;

    // 索引对应象限：0=左下，1=右下，2=左上，3=右上
    [
        (x_min, y_min, x_mid, y_mid),
        (x_mid, y_min, x_max, y_mid),
        (x_min, y_mid, x_mid, y_max),
        (x_mid, y_mid, x_max, y_max)
    ]
}

impl<'a> QuadBranch<'a> {
    fn new(rect: (f32, f32, f32, f32)) -> Self {
        let rects = divide_into_4(rect);
        QuadBranch {
            cells: [
                QuadTwig::new(rects[0]),
                QuadTwig::new(rects[1]),
                QuadTwig::new(rects[2]),
                QuadTwig::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
    }
}

impl<'a> QuadTree<'a> {
    fn new(rect: (f32, f32, f32, f32)) -> Self {
        let rects = divide_into_4(rect);
        QuadTree {
            cells: [
                QuadBranch::new(rects[0]),
                QuadBranch::new(rects[1]),
                QuadBranch::new(rects[2]),
                QuadBranch::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
    }

    fn into_point_and_leaf_iter(&self) -> PointAndLeafIterator<'_> {
        PointAndLeafIterator {
            ptr: self,
            index: (0, 0, 0, 0)
        }
    }

    // 新增迭代器构造方法
    fn into_point_and_leaf_and_neighbours_iter(&self) -> PointAndLeafAndNeighboursIterator<'_> {
        PointAndLeafAndNeighboursIterator {
            ptr: self,
            index: (0, 0, 0, 0)
        }
    }
}

2. 实现PointAndLeafAndNeighboursIterator #

迭代器核心逻辑：沿用PointAndLeafIterator的遍历逻辑逐个访问点，对每个点所在叶子计算周围9个方向的相邻叶子，超出QuadTree范围则返回空Vec。

/// 可返回每个点、其所在叶子及相邻叶子的迭代器
struct PointAndLeafAndNeighboursIterator<'a> {
    ptr: &'a QuadTree<'a>,
    index: (usize, usize, usize, usize)
}

impl<'a> Iterator for PointAndLeafAndNeighboursIterator<'a> {
    /// 返回点及周围9个叶子（无对应叶子时返回空叶子）
    type Item = (&'a (f32, f32), [Vec<&'a (f32, f32)>; 9]);

    /// 从(0,0,0,0)开始遍历，直至所有点遍历完成
    fn next(&mut self) -> Option<Self::Item> {
        let (branch_idx, twig_idx, leaf_idx, point_idx) = &mut self.index;

        // 获取当前层级的叶子与点
        let branch = &self.ptr.cells[*branch_idx];
        let twig = &branch.cells[*twig_idx];
        let leaf = &twig.cells[*leaf_idx];
        let point = leaf.vec.get(*point_idx);

        if let Some(point) = point {
            // 计算当前叶子的9个相邻叶子
            let neighbours = self.get_neighbour_leaves(*branch_idx, *twig_idx, *leaf_idx);
            *point_idx += 1;
            return Some((point, neighbours));
        }

        // 切换索引：当前叶子遍历完 → 下一个叶子
        *point_idx = 0;
        *leaf_idx += 1;
        if *leaf_idx < 4 {
            return self.next();
        }

        // 当前枝节点遍历完 → 下一个枝节点
        *leaf_idx = 0;
        *twig_idx += 1;
        if *twig_idx < 4 {
            return self.next();
        }

        // 当前分支节点遍历完 → 下一个分支节点
        *twig_idx = 0;
        *branch_idx += 1;
        if *branch_idx < 4 {
            return self.next();
        }

        // 遍历完成
        None
    }
}

impl<'a> PointAndLeafAndNeighboursIterator<'a> {
    // 将索引转换为象限坐标（x:0/1, y:0/1）
    fn idx_to_xy(idx: usize) -> (usize, usize) {
        match idx {
            0 => (0, 0),
            1 => (1, 0),
            2 => (0, 1),
            3 => (1, 1),
            _ => panic!("Invalid index: {}", idx),
        }
    }

    // 将象限坐标转换为索引
    fn xy_to_idx(x: usize, y: usize) -> usize {
        match (x, y) {
            (0, 0) => 0,
            (1, 0) => 1,
            (0, 1) => 2,
            (1, 1) => 3,
            _ => panic!("Invalid coordinates: ({}, {})", x, y),
        }
    }

    // 获取当前叶子周围9个相邻叶子的点集合
    fn get_neighbour_leaves(&self, branch_idx: usize, twig_idx: usize, leaf_idx: usize) -> [Vec<&'a (f32, f32)>; 9] {
        // 9个方向：左上、上、右上、左、中、右、左下、下、右下
        let directions = [(-1, -1), (0, -1), (1, -1),
                          (-1, 0),  (0, 0),  (1, 0),
                          (-1, 1),  (0, 1),  (1, 1)];

        let mut neighbours = [Vec::new(); 9];

        for (i, &(dx, dy)) in directions.iter().enumerate() {
            // 转换当前索引为坐标
            let (branch_x, branch_y) = Self::idx_to_xy(branch_idx);
            let (twig_x, twig_y) = Self::idx_to_xy(twig_idx);
            let (leaf_x, leaf_y) = Self::idx_to_xy(leaf_idx);

            // 计算偏移后的坐标，逐层向上调整
            let mut new_leaf_x = leaf_x as i32 + dx;
            let mut new_leaf_y = leaf_y as i32 + dy;
            let mut new_twig_x = twig_x as i32;
            let mut new_twig_y = twig_y as i32;
            let mut new_branch_x = branch_x as i32;
            let mut new_branch_y = branch_y as i32;

            // 调整叶子坐标，超出则移动枝节点
            if new_leaf_x < 0 {
                new_leaf_x = 1;
                new_twig_x -= 1;
            } else if new_leaf_x > 1 {
                new_leaf_x = 0;
                new_twig_x += 1;
            }

            if new_leaf_y < 0 {
                new_leaf_y = 1;
                new_twig_y -= 1;
            } else if new_leaf_y > 1 {
                new_leaf_y = 0;
                new_twig_y += 1;
            }

            // 调整枝节点坐标，超出则移动分支节点
            if new_twig_x < 0 {
                new_twig_x = 1;
                new_branch_x -= 1;
            } else if new_twig_x > 1 {
                new_twig_x = 0;
                new_branch_x += 1;
            }

            if new_twig_y < 0 {
                new_twig_y = 1;
                new_branch_y -= 1;
            } else if new_twig_y > 1 {
                new_twig_y = 0;
                new_branch_y += 1;
            }

            // 检查分支节点是否在QuadTree范围内
            if new_branch_x < 0 || new_branch_x > 1 || new_branch_y < 0 || new_branch_y > 1 {
                continue;
            }

            // 转换回索引并获取叶子
            let new_branch_idx = Self::xy_to_idx(new_branch_x as usize, new_branch_y as usize);
            let new_twig_idx = Self::xy_to_idx(new_twig_x as usize, new_twig_y as usize);
            let new_leaf_idx = Self::xy_to_idx(new_leaf_x as usize, new_leaf_y as usize);

            let branch = &self.ptr.cells[new_branch_idx];
            let twig = &branch.cells[new_twig_idx];
            let leaf = &twig.cells[new_leaf_idx];
            neighbours[i] = leaf.vec.clone();
        }

        neighbours
    }
}

3. 使用示例 #

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, neighbours) in quadtree.into_point_and_leaf_and_neighbours_iter() {
        println!("Point: {:?}", point);
        println!("Neighbour leaves (9 total):");
        for (i, leaf_points) in neighbours.iter().enumerate() {
            println!("  Direction {}: {:?}", i, leaf_points);
        }
    }
}

关键逻辑说明 #

• 索引与坐标转换：将每个层级的4个节点索引转换为(0/1, 0/1)的坐标，方便计算相邻位置
• 逐层偏移调整：当叶子坐标超出当前枝节点范围时，调整枝节点坐标；枝节点超出分支节点范围时，调整分支节点坐标，最终判断是否在QuadTree范围内
• 边界处理：超出QuadTree范围的方向直接返回空Vec，保证返回的数组始终有9个元素

内容的提问来源于stack exchange，提问作者Blue7