vlambda博客
学习文章列表

作者丨大魔王

来源丨小小算法

使用栈实现二叉树的遍历

我们遍历二叉树的目的是什么？目的是对二叉树的结点进行操作。

一般情况下，我们只是对一棵树的所有结点进行操作，前序、中序、还是后序遍历都无所谓，但涉及到对结点操作顺序有严格要求的问题时，我们就必须考虑使用哪种遍历方式来解决题目要求中隐藏的顺序问题。

使用栈来控制对结点的操作顺序时，我们只需要弄清楚两个问题：

1. 结点何时进栈？
2. 结点何时出栈？

其实很简单，下面我们结合这个简单二叉树模型，一起来解决上面这两个问题。

为了统一三种遍历方式的算法思想，在中序和后序遍历中，将会引入 set 作为辅助数据结构。

先序遍历

结点何时进栈呢？

先序遍历的遍历顺序是 node, left, right . 所以当 node 为栈顶元素时，node 先出栈，right 和 left 再依次进栈. 需要注意栈是后进先出的数据结构，所以是 right 先进 left 后进。

结点何时出栈呢？

只要该结点满足是栈顶元素即可。

代码

void preorderTraversal(TreeNode* root) {
    vector<int> ret;
    stack<TreeNode*> s;
    s.push(root);
    while (!s.empty()) {
        TreeNode* node = s.top();
        // 栈顶元素出栈，之后right和left依次入栈
        s.pop();
        if (node == NULL) {
            continue;
        }
        cout << node->val << endl;
        s.push(node->right);
        s.push(node->left);
    }
}

中序遍历

结点何时进栈呢？

中序遍历的操作顺序是 left, node, right . 因此当 node 为栈顶元素时，left 直接进栈 , right 要等 node 出栈后再进栈。

结点何时出栈呢？

该结点必须是栈顶元素，并且它的左结点 left 已经操作完了（也就是出栈了），它才可以进栈。

这里有一个难题，node 怎么知道 left 是否已经出栈了呢？同样 right 怎么知道 node 是否已经出栈了呢？

我们可以使用一个 set 或者 map 来保存已经出栈的结点，如果 left 在 set 里就说明它已经出栈了。

代码

void inorderTraversal(TreeNode* root) {
    stack<TreeNode*> s;
    s.push(root);
    set<TreeNode*> isOut;
    isOut.insert(NULL);
    while (!s.empty()) {
        TreeNode* node = s.top();
        if (node == NULL) {
            s.pop();
            continue;
        }
        // 如果左结点已经操作完，输出当前结点，并将它加入 set
        if (isOut.find(node->left) != isOut.end()) {
            cout << node->val << endl;
            // 栈顶元素出栈，右结点进栈
            s.pop();
            isOut.insert(node);
            s.push(node->right);
        // 否则左结点入栈
        } else {
            s.push(node->left);
        }
    }
}

后序遍历

结点何时进栈呢？

后序遍历的操作顺序是 left, right, node . 当 node 为栈顶元素时，因此进栈顺序是 right , left .

结点何时出栈呢？

首先它必须是栈顶元素，然后它的 left 和 right 都已经出栈了，它才可以出栈。

判断 left 或 right 是否已经出栈，同样是使用 set 或 map .

代码

vector<int> postorderTraversal(TreeNode* root) {
    set<TreeNode*> isOut;
    isOut.insert(NULL);
    stack<TreeNode*> s;
    s.push(root);
    while (!s.empty()) {
        TreeNode* node = s.top();
        if (node == NULL) {
            s.pop();
            continue;
        }
        // 如果左右结点都已操作完，输出当前结点，并将它加入 set
        if (isOut.find(node->left) != isOut.end() && isOut.find(node->right) != isOut.end()) {
            cout << node->val << endl;
            s.pop();
            isOut.insert(node);
        }
        // 如果右结点还未操作，右结点入栈
        if (isOut.find(node->right) == isOut.end()) {
            s.push(node->right);
        }
        // 如果左结点还未操作，左结点入栈
        if (isOut.find(node->left) == isOut.end()) {
            s.push(node->left);
        }
    }
}

最后

这样我们就通过引入 set 作为辅助数据结构，统一了遍历思想，轻松的实现了使用栈遍历二叉树的三种遍历方式。

感谢您的观看，如果本文对您有帮助，不妨点个在看支持一下，谢谢！