c++-二叉树遍历for loop统一迭代法

序

二叉树遍历是个老生常谈的问题，下面提供一种统一的for loop 迭代法
修改便利方式
PRE 先序遍历 IN 中序遍历 POST 后序遍历
支持hook

不同的遍历方式只需要修改seq的值

leetcode 原题实战

94. 二叉树的中序遍历

https://leetcode.cn/problems/binary-tree-inorder-traversal/

... 
treeAdaptor.h 代码
...
class Solution {
public:
    vector<int> inorderTraversal(TreeNode* root) {
        return getSeqValue((TreeNodeAdapter*)root, IN);
    }
};

98 验证二叉搜索树

https://leetcode.cn/problems/validate-binary-search-tree/

class Solution {
public:
    bool isValidBST(TreeNode* root) {
        if (nullptr == root) {
            return true;
        }
        long long pre = INT64_MIN;
        setTraversal(IN);
        for (auto v : *(TreeNodeAdapter*)root) {
            if (v <= pre) {
                return false;
            }
            pre = v;
        }
        return true;
    }
};

99. 恢复二叉搜索树

https://leetcode.cn/problems/recover-binary-search-tree/

// 实现hook 
struct AnsObj : public Hook{
    TreeNodeAdapter* pre{};
    TreeNodeAdapter* left{};
    TreeNodeAdapter* right{};
	// 逻辑代码
    bool operator()(TreeNodeAdapter* node) override{
        if (nullptr == node) {
            return true;
        }
        if (pre && pre->val > node->val) {
            right = node;
            if (nullptr == left) {
                left = pre;
            } else {
                return true;
            }
        }
        pre = node;
        return false;
    }
	
	// 遍历结束后的操作
    void after() override {
        swap(left->val, right->val);
    }
};

class Solution {
public:
    void recoverTree(TreeNode* root) {
        if (nullptr == root) {
            return;
        }
        AnsObj ans;
        traversal((TreeNodeAdapter*)root, IN, &ans);
    }
};

144. 二叉树的前序遍历

https://leetcode.cn/problems/binary-tree-preorder-traversal/

class Solution {
public:
    vector<int> preorderTraversal(TreeNode* root) {
        return getSeqValue((TreeNodeAdapter*)root, PRE);
    }
};

145. 二叉树的后序遍历

https://leetcode.cn/problems/binary-tree-postorder-traversal/

treeAdaptor.h 代码
...
class Solution {
public:
    vector<int> postorderTraversal(TreeNode* root) {
        return getSeqValue((TreeNodeAdapter*)root, POST);
    }
};

源码

treeAdaptor.h

treeAdaptor.h

//      treeAdaptor.h
#ifndef UNTITLED3_TREE_H
#define UNTITLED3_TREE_H
#include <iostream>
#include <vector>
#include <stack>

using namespace std;
using VALUE_TYPE  = int;

struct TreeNodeIterator;
enum SEQ{
    PRE = 0,
    IN = 1,
    POST = 2,
};

static SEQ seq = POST;

void setTraversal(int v) {
    if (v == 0) {
        seq = PRE;
    } else if (v == 1) {
        seq = IN;
    } else if (v == 2){
        seq = POST;
    } else {
        seq = PRE;
    }
}


struct TreeNodeAdapter {
    explicit TreeNodeAdapter(const VALUE_TYPE &val, TreeNodeAdapter *left = nullptr, TreeNodeAdapter *right = nullptr) : val(val), left(left), right(right) {}

    TreeNodeIterator begin();

    static TreeNodeIterator end();

    VALUE_TYPE val{};
    struct TreeNodeAdapter *left{}, *right{};

};

struct Hook {
    virtual bool operator()(TreeNodeAdapter* node) {
        return false;
    }
    virtual void after() {

    }
    bool done{};

    bool isDone() const{
        return done;
    }
};

struct EmptyHook : public Hook {};
EmptyHook emptyHook;

static Hook* hook = &emptyHook;

struct TreeNodeIteratorImpAbs {
    stack<TreeNodeAdapter *> st;
    TreeNodeAdapter* now{};

    virtual void next(TreeNodeAdapter *node) = 0;

    explicit TreeNodeIteratorImpAbs(TreeNodeAdapter *root) {
        if (root) {
            st.push(root);
        }
    }

    bool getNext() {
        while (!hook->isDone() && !st.empty()) {
            // 当前指针指向栈顶节点
            TreeNodeAdapter *cur = st.top();
            if (cur) {    // 当前节点不为空
                st.pop();
                // TODO:按照各自遍历顺序将节点入栈，入栈中间节点后要入栈一个空节点
                next(cur);
            } else {  //当前节点为空
                // 弹出空节点
                st.pop();
                // 弹出中间节点并将它的值加入结果数组
                now = st.top();
                st.pop();
                if ((*hook)(now)) {
                    hook->done = true;
                }
                return true;
            }
        }
        now = nullptr;
        return false;
    }
};

struct preOrderTreeNodeIteratorImp : public TreeNodeIteratorImpAbs {
    explicit preOrderTreeNodeIteratorImp(TreeNodeAdapter *root) : TreeNodeIteratorImpAbs(root) {}

    void next(TreeNodeAdapter *node) override {
        if (node->right) {
            st.push(node->right);
        }
        if (node->left) {
            st.push(node->left);
        }
        st.push(node);
        st.push(nullptr);
    }
};

struct inOrderTreeNodeIteratorImp : public TreeNodeIteratorImpAbs {
    explicit inOrderTreeNodeIteratorImp(TreeNodeAdapter *root) : TreeNodeIteratorImpAbs(root) {}

    void next(TreeNodeAdapter *node) override {
        if (node->right) {
            st.push(node->right);
        }
        st.push(node);
        st.push(nullptr);
        if (node->left) {
            st.push(node->left);
        }
    }
};

struct postOrderTreeNodeIteratorImp : public TreeNodeIteratorImpAbs {
    explicit postOrderTreeNodeIteratorImp(TreeNodeAdapter *root) : TreeNodeIteratorImpAbs(root) {}

    void next(TreeNodeAdapter *node) override {
        st.push(node);
        st.push(nullptr);
        if (node->right) {
            st.push(node->right);
        }
        if (node->left) {
            st.push(node->left);
        }
    }
};

struct TreeNodeIterator {
    TreeNodeIteratorImpAbs *node;

    explicit TreeNodeIterator(TreeNodeIteratorImpAbs *node) : node(node) {
        node->getNext();
    }

    int &operator*() {
        return node->now->val;
    }

    TreeNodeIterator operator++() {
        node->getNext();
        return *this;
    }

    bool operator!=(TreeNodeIterator other) const {
        return node->now != other.node->now;
    }
};

TreeNodeIterator TreeNodeAdapter::begin() {
    if (seq == PRE) {
        return TreeNodeIterator(new preOrderTreeNodeIteratorImp(this));
    }
    if (seq == IN) {
        return TreeNodeIterator(new inOrderTreeNodeIteratorImp(this));
    }
    return TreeNodeIterator(new postOrderTreeNodeIteratorImp(this));
}

TreeNodeIterator TreeNodeAdapter::end() {
    return TreeNodeIterator(new preOrderTreeNodeIteratorImp(nullptr));
}

       对外层

void traversal(TreeNodeAdapter* root, SEQ se, Hook* func = nullptr) {
    setTraversal(se);
    if (func) {
        hook = func;
    }
    auto beg = ((TreeNodeAdapter*)root)->begin();
    auto end = ((TreeNodeAdapter*)root)->end();
    for (; beg != end; ++beg);
    hook->after();
}

vector<TreeNodeAdapter*> getSeqTreeNode(TreeNodeAdapter* root, SEQ se = PRE) {
    setTraversal(se);
    vector<TreeNodeAdapter*> data;
    auto beg = ((TreeNodeAdapter*)root)->begin();
    auto end = ((TreeNodeAdapter*)root)->end();
    for (; beg != end; ++beg) {
        data.push_back(beg.node->now);
    }
    return data;
}

vector<int> getSeqValue(TreeNodeAdapter* root, SEQ se = PRE) {
    setTraversal(se);
    vector<int> data;
    if (!root) {
        return data;
    }
    for (auto v : *root) {
        data.push_back(v);
    }
    return data;
}

#endif //UNTITLED3_TREE_H

main.cpp

#include "tree.h"

// 这个是我们遇见的可能的TreeNode的实现方式，
// 我们需要保持TreeNodeAdapter 和 TreeNode 的在内存布局是一样的（成员一样且声明顺序一致）

struct TreeNode {
    explicit TreeNode(const VALUE_TYPE &val, TreeNode *left = nullptr, TreeNode *right = nullptr) : val(val), left(left), right(right) {}

    struct TreeNode *left{}, *right{};
    VALUE_TYPE val{};

};

// 生成一颗简单的树
TreeNodeAdapter* makeTree() {
    auto* root = new TreeNode(0);
    root->left = new TreeNode(1);
    root->right = new TreeNode(2);
    root->left->right = new TreeNode(5);
    root->right->left = new TreeNode(4);
    root->left->right->left = new TreeNode(3);
    return reinterpret_cast<TreeNodeAdapter *>(root);
}

int main() {
	// 默认为前序遍历
    setTraversal(PRE);
    cout << "前序遍历：";
    for (auto r : *makeTree()) {
        cout << r << " ";
    }
    cout << endl;
    
    setTraversal(IN);
    cout << "中序遍历：";
    for (auto r : *makeTree()) {
        cout << r << " ";
    }
    cout << endl;
    
    setTraversal(POST);
    cout << "后序遍历：";
    for (auto r : *makeTree()) {
        cout << r << " ";
    }
    cout << endl;
}

输出

c+±二叉树遍历for loop统一迭代法
c+±二叉树遍历for loop统一迭代法
c+±二叉树遍历for loop统一迭代法
c+±二叉树遍历for loop统一迭代法
c+±二叉树遍历for loop统一迭代法
c+±二叉树遍历for loop统一迭代法
c+±二叉树遍历for loop统一迭代法
c+±二叉树遍历for loop统一迭代法
c+±二叉树遍历for loop统一迭代法
c+±二叉树遍历for loop统一迭代法