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

#include 
#include 
#include 
#include 

using namespace std;

class Visitor;

class TreeNode;

class Visitor {
public:
    virtual vector<int> visit(TreeNode *root) = 0;
};

template<typename Derive>
class VisitingBase {
public:
    vector<int> accept(Visitor *v) {
        return v->visit((Derive *) (this));
    }
};

class TreeNode : public VisitingBase<TreeNode> {
public:
    int val;
    struct TreeNode *left, *right;

    explicit TreeNode(int val = 0, TreeNode *left = nullptr, TreeNode *right = nullptr) : val(val), left(left),
                                                                                          right(right) {}
};

class TreeNodeTraversalHelper {
protected:
    vector<int> ans;
    int startIndex;
    TreeNode* cur = nullptr;
    TreeNode* TreeNode::* child[2] = {&TreeNode::left, &TreeNode::right};

    void push_back() {
        assert(NULL != cur);
        ans.push_back(cur->val);
    }

    void reverse() {
        ::reverse(ans.begin(), ans.end());
    }

    virtual void PreMeetHook() {}
    virtual void InMeetHook() {}
    virtual void AfterVisitHook() {};

    int getPreIndex() const {
        return startIndex & 0x1;
    }

    int getInIndex() const {
        return (startIndex + 1) & 0x1;
    }

public:
    TreeNodeTraversalHelper(int start) : startIndex(start) {}

    vector<int> doVisit(TreeNode* root) {
        if (!root) {
            return {};
        }
        stack<TreeNode*> st;
        cur = root;
        while (cur || !st.empty()) {
            while (cur) {
                PreMeetHook();
                st.push(cur);
                cur = cur->*child[getPreIndex()];
            }
            cur = st.top();
            st.pop();
            InMeetHook();
            cur = cur->*child[getInIndex()];
        }
        AfterVisitHook();
        return ans;
    }
};


class PreVisitor : public Visitor, public TreeNodeTraversalHelper{
private:
    void PreMeetHook() override {
        push_back();
    }

public:
    PreVisitor() : TreeNodeTraversalHelper(0){}

    vector<int> visit(TreeNode *root) override {
        return doVisit(root);
    }
};

class InVisitor : public Visitor, public TreeNodeTraversalHelper{
protected:
    void InMeetHook() override {
        push_back();
    }

public:
    InVisitor() : TreeNodeTraversalHelper(0) {}

    vector<int> visit(TreeNode *root) override {
        return doVisit(root);
    }
};

class PostVisitor : public Visitor, public TreeNodeTraversalHelper{
protected:
    void PreMeetHook() override {
        push_back();
    }

    void AfterVisitHook() override {
        reverse();
    }

public:
    PostVisitor():TreeNodeTraversalHelper(1) {}

    vector<int> visit(TreeNode *root) override {
        return doVisit(root);
    }
};

TreeNode *makeTree() {
    TreeNode *root = new TreeNode(0);
    root->left = new TreeNode(1);
    root->right = new TreeNode(2);
    root->left->left = new TreeNode(3);
    root->right->left = new TreeNode(6);
    root->right->right = new TreeNode(4);
    root->right->right->right = new TreeNode(5);
    return root;
}


int main() {
    TreeNode *root = makeTree();
    PreVisitor pr;
    for (auto v: root->accept(&pr)) {
        cout << v << "\t";
    }
    cout << endl;

    InVisitor iv;
    for (auto v: root->accept(&iv)) {
        cout << v << "\t";
    }
    cout << endl;

    PostVisitor pv;
    for (auto v: root->accept(&pv)) {
        cout << v << "\t";
    }
    cout << endl;
}

你可能感兴趣的:(c++,开发语言)