二叉树的遍历-栈 - BUG研发中心

一个在数据结构领域被艹了又艹的话题。本文只涉及使用栈来遍历二叉树（不包含层次遍历），并将其包装为迭代器。

UPDATE 2016/06/22: 将栈的部分从原来的大文章中分离出来

UPDATE 2016/10/28: 拿回了属于自己的迭代器

直接遍历

有关约定

树节点定义如下：

struct TreeNode {
    int data;
    TreeNode *left,*right;
    
    TreeNode(int x):data(x),left(nullptr),right(nullptr) {}
};

代码中的Func是函数类，也可以用函数指针/std::function替换之。

使用系统栈

这绝对是简单到爆的一个实现方法了，不可能还有比这更简单的了。几乎就是照着这几种遍历的定义翻译到程序代码，因此不做过多解释。

前序遍历

template <class Func> void preorder(TreeNode *root,Func visit){
    if(root==nullptr)
        return;
    visit(root->data);
    preorder(root->left,visit);
    preorder(root->right,visit);
}

中序遍历

template <class Func> void inorder(TreeNode *root,Func visit){
    if(root==nullptr)
        return;
    inorder(root->left,visit);
    visit(root->data);
    inorder(root->right,visit);
}

后序遍历

template <class Func> void postorder(TreeNode *root,Func visit){
    if(root==nullptr)
        return;
    postorder(root->left,visit);
    postorder(root->right,visit);
    visit(root->data);
}

使用自定义栈

难度开始逐步提升了。在抛开了系统栈之后，我们应该如何使用自定义栈来替换掉系统栈呢？

前序遍历

观察原来使用系统栈的算法，先访问当前节点，之后向左走，再之后向右走。如果只是对系统栈无脑地模拟的话，向左走时把当前节点入栈，完成后退栈；向右走时同法。但在向右走时这一个调用（preorder(root->right,visit)）属于尾递归，此时不需要入栈，直接替换掉对应变量值即可。

总结一下，先访问当前节点；将右边入栈；节点值替换成左节点，循环。当前支路走无可走之时，弹栈，继续；此时如果栈空，结束。循环这整个过程。

#include <stack>

template <class Func> void preorder(TreeNode *root,Func visit){
    if(root==nullptr)
        return;
    std::stack<TreeNode*> stack;
    while(root!=nullptr){
        visit(root->data);
        if(root->right!=nullptr)
            stack.push(root->right);
        if(root->left!=nullptr)
            root=root->left;
        else{
            if(stack.empty())
                return;
            root=stack.top();
            stack.pop();
        }
    }
}

中序遍历

同样对系统栈进行模拟，参考前序的流程，中序的流程如下：将当前节点入栈，向左走走到黑；不可走时弹栈（如果栈空则结束），访问该节点，之后向右走一步。如果向右走遇到空，重复弹栈-访问-向右的操作。循环这整个过程。

#include <stack>

template <class Func> void inorder(TreeNode *root,Func visit){
    if(root==nullptr)
        return;
    std::stack<TreeNode*> stack;
    while(root!=nullptr||!stack.empty()){
        if(root!=nullptr){
            stack.push(root);
            root=root->left;
        } else {
            if(stack.empty())
                return;
            root=stack.top();
            stack.pop();
            visit(root->data);
            root=root->right;
        }
    }
}

为简化代码起见，把“弹栈-访问-向右”的循环合并到主流程里面了。

后序遍历

后序遍历就有些复杂了（Leetcode上非递归前序，中序都是Medium, 后序就成了Hard我会说？）。

将当前节点入栈，向左走；不可走时弹栈，向右走，跳回前一步；不可走时弹栈，栈空时结束。问题在于，此时弹栈之后，可能是要向左走，也可能是要向右走，而且同一个节点可能会被弹多次。考虑下面这棵树：

在系统栈作用下，在4不可走时，节点2被弹出，随后又入栈，走向5, 此后再一次弹出2, 这次是真正的弹出了。因此，重点在于判断节点被弹出后，是否还要向右走。主要是需要判断是否是第二次弹出，因此可设一个visit数组来记录访问状态，但太浪费空间。

这时就要用到后序遍历的特性了。记当前节点为node, 若node被访问，则其左右子树已经遍历完毕，应该向上回溯了。如果node为parent的左子树，代表着parent的左子树遍历完成，接下来就是右子树，应该向parent的右边进发；如果node是parent的右子树，代表parent的左右子树已经遍历完毕，这时它的情况和node一样，循环进行处理。

我们设置一个辅助变量last(初始设为空), 用来记录上次访问的节点。遍历时，节点一路向左一路入栈走到成空为止，此后取栈顶parent, 如果上次访问的last不为parent的右子树，且右子树不为空，这时要向右进发；其他情况（即parent的左右子树已经遍历完毕）就该访问parent了。之后，更新last记录。

#include <stack>

template <class Func> void postorder(TreeNode *root,Func visit){
    std::stack<TreeNode*> stack;
    TreeNode *last=nullptr;
    while(root!=nullptr||!stack.empty()){
        if(root!=nullptr){
            stack.push(root);
            root=root->left;
        } else {
            TreeNode *parent=stack.top();
            if(last!=parent->right && parent->right!=nullptr)
                root=parent->right;
            else {
                visit(parent->data);
                last=parent;
                stack.pop();
            }
        }
    }
}

仙术

此方法根据树的遍历定义得来。树的前序遍历是“根-左-右”，后序遍历是“左-右-根”。修改前序遍历的代码，使其顺序变为“根-右-左”，再将修改后得到的前序遍历序列逆置，即可得到后序遍历的“左-右-根”顺序。由于该方法需要一个额外的记录序列的容器，Leetcode的题目正好就是要获得一个这样的序列，因此该序列容器不占空间复杂度，在我这个遍历算法中就不一样了。解法放在这里，给各位扩展一下思路。

该方法来自Leetcode讨论区的这篇帖子。

迭代器遍历

有关约定

迭代器有关接口声明如下：

#include <stack>

class iterator{
    public:
        iterator(TreeNode *root); // 待实现
        iterator(const iterator &o):m_cursor(o.m_cursor) {}
        
        int& operator*() const {return m_cursor.top()->data;}
        int* operator->() const {return &operator*();}
        iterator& operator++(); // 待实现
        iterator operator++(int) {auto i=*this; operator++(); return i;}
    private:
        std::stack<TreeNode*> m_cursor;
};

operator++()所实现的操作是找到在遍历序列中，当前节点的下一个节点。我们应该认为当前节点已经访问过了。
使用栈时的迭代器内部封装的是一个栈，复制起来开销太大，你可以在此基础实现上移动语义以降低开销。
在下文中，虽然三种遍历算法的迭代器都叫iterator, 但实际上是三种不同的类型。为描述简单起见我共用一个名称。
在这里operator->()是非必须的，如果使用会导致编译错误。但设计成泛型迭代器的话那么就是必须要重载的一个重要运算符了。
如果想和STL接轨，迭代器实现起来不会如此简单，还要定义一大堆类型别名，设置迭代器种类等。该迭代器由于只能自增，因此属于ForwardIterator.

具体实现

主算法已经在上面讨论过了，因此在这里主要是将循环的单步操作提取出来作为迭代器的自增操作。初始化迭代器时，应该找到这种遍历方式生成的序列中的第一个节点。

前序遍历

前序的第一个节点就是根。

iterator::iterator(TreeNode *root):m_cursor() {
    if(root!=nullptr)
        m_cursor.push(root);
}

iterator& iterator::operator++() {
    if(m_cursor.empty())
        return *this;
    auto tmp=m_cursor.top();
    m_cursor.pop();
    if(tmp->right != nullptr)
        m_cursor.push(tmp->right);
    if(tmp->left != nullptr)
        m_cursor.push(tmp->left);
    return *this;
}

中序遍历

中序遍历的第一个节点是最左的一个节点。

iterator::iterator(TreeNode *root):m_cursor() {
    while(root!=nullptr){
        m_cursor.push(root);
        root=root->left;
    }
}

iterator& iterator::operator++() {
    if(m_cursor.empty())
        return *this;
    auto ptr = m_cursor.top()->right;
    m_cursor.pop();
    while(ptr != nullptr) {
        m_cursor.push(ptr);
        ptr = ptr->left;
    }
}

后序遍历

后序遍历的第一个节点沿左边的支路向下的第一个叶子节点。

为了不给后序迭代器增加一个last成员，将上面的遍历算法做出一点修改：

// 辅助用函数
void next(TreeNode *root,std::stack<TreeNode*> &s){
    for(;;){
        s.push(root);
        if(root->left != nullptr)
            root = root->left;
        else if(root->right != nullptr)
            root = root->right;
        else
            return;
    }
}

template <class Func> void postorder(TreeNode *root,Func visit){
    std::stack<TreeNode*> stack;
    next(root,stack);
    while(!stack.empty()){
        root=stack.top();stack.pop();
        visit(root->data);
        if(stack.empty())
            return;
        TreeNode *parent=stack.top(), *tmp=parent->right;
        if(root==parent->left && tmp!=nullptr)
            next(tmp,stack);
    }
}

迭代器算法如下：

// 辅助用函数
void next(TreeNode *root,std::stack<TreeNode*> &s){
    for(;;){
        s.push(root);
        if(root->left != nullptr)
            root = root->left;
        else if(root->right != nullptr)
            root = root->right;
        else
            return;
    }
}

iterator::iterator(TreeNode *root):m_cursor() {
    if(root==nullptr)
        return;
    next(root,m_cursor);
}


iterator& iterator::operator++() {
    if(m_cursor.empty())
        return *this;
    auto ptr=m_cursor.top();
    m_cursor.pop_back();
    if(m_cursor.empty())
        return *this;
    auto parent = m_cursor.top(), tmp = parent->right;
    if(ptr == parent->left && tmp != nullptr)
        next(tmp,m_cursor);
    return *this;
}