现代C++如何为KDTree实现支持范围for循环的迭代器？

首先先指出你现有遍历代码的两个明显bug：

• 现有traverseNode实现漏处理根节点：你调用traverseNode(root, func)，但函数内部只对node->left和node->right调用func，根节点本身不会被回调
• 右子树递归分支逻辑错误：traverseNode(node->left, func)应该改为traverseNode(node->right, func)，否则右子树完全不会被遍历，还会重复遍历左子树

实现切入点 #

C++范围for语法的本质是编译器自动调用类的begin()和end()方法获取迭代器，之后循环执行*it、++it直到it == end()。你只需要完成两个核心改造即可：

1. 定义符合要求的迭代器类型，内部用栈维护树的遍历状态，完全避免递归开销和节点拷贝开销
2. 给KDTree类补充begin()和end()公共方法，返回对应的迭代器实例

完整实现代码 #

#include <stack>
#include <iterator>

struct Node {
    int id;
    Vector point;
    Node *left, *right;
    Node(int id, double x, double y) : id(id), point(x, y), left(nullptr), right(nullptr) {}
};

class KDTree {
    Node *root;
    void insertNode(Node *node, int id, double x, double y);
    void searchNode(Node *node, double x, double y, double r, vector<int> &ids);
    void removeNode(Node *node, int id);
    void clearNode(Node *node);
    void traverseNode(Node *node, function<void(Node*)> func);
public:
    KDTree() : root(nullptr) {}
    void remove(int id);
    void insert(int id, double x, double y);
    vector<int> search(double x, double y, double r);
    void clear();
    
    vector<Vector> traverse();
    void traverse(function<void(Node*)> func);

    // 新增前序遍历迭代器，和常规树遍历逻辑对齐
    class Iterator {
        std::stack<Node*> path;
    public:
        using value_type = Node*;
        using difference_type = std::ptrdiff_t;
        using pointer = Node**;
        using reference = Node*&;
        using iterator_category = std::forward_iterator_tag;

        Iterator() = default;
        explicit Iterator(Node* root) {
            if (root) path.push(root);
        }

        reference operator*() {
            return path.top();
        }

        Iterator& operator++() {
            Node* curr = path.top();
            path.pop();
            // 先压右节点再压左节点，弹出时先左后右，符合前序遍历顺序
            if (curr->right) path.push(curr->right);
            if (curr->left) path.push(curr->left);
            return *this;
        }

        bool operator!=(const Iterator& other) const {
            return path != other.path;
        }
    };

    Iterator begin() {
        return Iterator(root);
    }

    Iterator end() {
        return Iterator();
    }
};

效果说明 #

改造完成后你可以直接用以下语法遍历所有节点：

for (auto* node : kdtree) {
    // 直接访问node的id、point等成员
}

该实现没有额外节点拷贝开销，也没有递归或函数回调的额外开销，遍历时间复杂度和递归遍历完全一致，额外空间复杂度为O(h)（h为树的高度），远低于返回vector方案的O(n)空间开销。如果需要其他遍历顺序（中序、后序、层序），只需要修改迭代器operator++的压栈逻辑即可。

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