二叉树的迭代遍历写法

在遍历二叉树时，对于当前节点有两种操作：

取值，即取出 node->val 的值，写入结果数组
访问左右子树

因此，二叉树的递归写法非常简单，对于前序、中序、后序遍历，只需要调整取值和访问左右子树操作的顺序即可。相比之下，用迭代方式遍历二叉树写起来就不那么直观。

我们都知道，递归实际上就是函数调用栈的增长，因此迭代遍历二叉树，本质上就是用一个栈来表达递归。那么递归遍历二叉树时，代码在执行层面具体发生了什么？我们以一棵最简单的二叉树为例：

首先是前序遍历，递归代码如下：

typedef struct TreeNode {
    int val;
    TreeNode* left;
    TreeNode* right;
    TreeNode() : val(0), left(nullptr), right(nullptr) {}
    TreeNode(int val_) : val(val_), left(nullptr), right(nullptr) {}
} TreeNode;

vector<int> res;

void preorderTraverse(TreeNode* node) {
    if (!node) return;              // END
    res.push_back(node->val);       // A
    preorderTraverse(node->left);   // B
    preorderTraverse(node->right);  // C
}

int main() {
    TreeNode* root = new TreeNode(1);
    // ...
    preorderTraverse(root);
    // output res
}

首次调用preorderTraverse()时，当前节点为 1，将 1 写入结果数组，即上面代码的 A 行。然后遍历左子树，即代码 B 行，此时递归调用preorderTraverse()，访问左子树，函数又从首行开始执行，到 A 行时，将 2 写入结果数组。然后节点 2 执行 B 、C 行，访问左右子树。由于它的左右子树都是空的，所以会直接返回。至此节点 2 函数执行完毕，也即节点 1 的左子树访问完毕，并且都写入了结果数组。所以，此时函数返回到节点 1 的 B 行，开始执行 C 行访问右子树。经历和左子树一样的过程后，返回到 C 行，遍历结束。

而对于中序遍历，情况变得有些不一样。中序遍历的代码如下：

void inorderTraverse(TreeNode* node) {
    if (!node) return;              // END
    inorderTraverse(node->left);    // B
    res.push_back(node->val);       // A
    inorderTraverse(node->right);   // C
}

当次调用inorderTraverse()时，当前节点为 1，但是我们现在先不取出 node->val 的值，而是先访问左子树，也就是去访问节点 2 。对于节点 2 ，我们也没有先取值，而是访问节点 2 的左子树，也就是空节点。空节点返回，那么节点 2 函数中的 B 行执行完毕，此时再去取节点 2 的值，写入结果数组。然后执行 C 行访问节点 2 的右子树，由于右子树为空，返回到 C 行，节点 2 为根的这棵树就遍历完毕了。此时，函数返回到节点 1 的 B 行，然后执行 A 行，取出节点 1 的值，写入结果数组。最后是右子树的遍历。

对于中序遍历，访问当前节点时，没有先使用 val ，而是使用了 left 指针，当左子树遍历结束返回后，才去取 val 的值，也就是说，“进入当前节点的函数”和“取出当前节点的值”这两个时刻中间还间隔了好长一段访问左子树的操作。所以，我们必须保存好当前节点的 val 值（以及 right 指针），才可以保证 B 行返回时可以继续操作。对于递归来说，这个状态的保存是自然而然的，当访问左子树时，实际上是函数调用栈的增长，而 B 行返回时，是增长的函数栈又被销毁，回到当前节点。而函数调用栈的每一帧都保存了当前的局部变量，以及代码执行的位置，所以返回到当前帧时，可以继续从 A 行往下执行。

不过现在，我们要用自己的栈来表达递归，那么麻烦就出现了。

首先明确，在迭代写法中，我们自己的栈里保存的是一个指针值TreeNode*，这样才可以拿到节点的所有变量。

对于前序遍历，我们可以很轻松地取出栈顶的TreeNode*并弹出，然后将 val 存到结果数组，再相继取出它的右、左子节点（这样弹出时就是左、右的顺序），如果不为空就压入栈。此后，这个TreeNode*就再也不会被用到了，因为对于 val 、left 、right 的操作已经全部做完了。

对于中序遍历，我们需要首先取出栈顶的TreeNode*（记为nodeX）并弹出，然后将右子树压入栈。这时问题来了，由于当前节点nodeX还没有取出 val ，所以我们势必不能直接丢掉当前节点，而是要让当前节点再次入栈，以便之后弹出时能够取它的 val 值。最后，还要将当前节点的左子树压入栈，以便左子树最先被弹出。可是，当代码继续执行，直到当前栈顶又为nodeX时，我们怎么知道nodeX的左右子树有没有入栈过呢？我们是直接取出 val 值写入结果数组，还是要让nodeX的左右子树入栈呢？所以，中序遍历时，栈里的节点就具有了两种状态（因为每个节点都会被压入栈两次），首次入栈并取出时，我们要将它的子树入栈，第二次入栈并取出时，我们要取它的 val 值。我们可以在节点第二次入栈时，紧接着压入一个 null 指针，来标识这个节点的状态，即下次取出一个 null 时，紧接着弹出的一个节点，只需要取它的 val 值就可以了。所以，这种迭代遍历的写法又叫标记法。

我们来尝试写出代码。

中序遍历的迭代写法：

vector<int> inorderIteration(TreeNode* root) {
    vector<int> res;
    stack<TreeNode*> stk;
    TreeNode* node = root;
    if(node) stk.push(node);

    while (!stk.empty()) {
        node = stk.top();
        stk.pop();              // 取出并弹出栈顶节点
        if (node != nullptr) {  // 如果不是null，则它是首次入栈，要将左右子节点压入栈
            if (node->right) stk.push(node->right);     // 右节点入栈
            stk.push(node);                             // 当前节点再次入栈
            stk.push(nullptr);                          // 压入一个null，表示当前节点是第二次入栈
            if (node->left) stk.push(node->left);       // 左节点入栈
        } else {                // 如果是null，则它是第二次入栈，需要取出 val 值
            node = stk.top();
            stk.pop();          // 再次弹出栈顶，取出真正的节点指针
            res.push_back(node->val);   // 取值并写入结果数组
        }
    }

    return res;
}

自然而然地，前序遍历和后序遍历也可以统一到这个框架下，仅仅需要改变左、中、右入栈的顺序，和递归写法一样。要注意，由于栈是先进后出，所以前序遍历的入栈顺序应该是右、左、中，后序遍历的入栈顺序应该是中、右、左。

前序遍历的迭代写法：

vector<int> preorderIteration(TreeNode* root) {
    vector<int> res;
    stack<TreeNode*> stk;
    TreeNode* node = root;
    if(node) stk.push(node);

    while (!stk.empty()) {
        node = stk.top();
        stk.pop();
        if (node != nullptr) {
            if (node->right) stk.push(node->right); // 右
            if (node->left) stk.push(node->left);   // 左
            stk.push(node);                         // 中
            stk.push(nullptr);
        } else {
            node = stk.top();
            stk.pop();
            res.push_back(node->val);
        }
    }

    return res;
}

后序遍历的迭代写法：

vector<int> postorderIteration(TreeNode* root) {
    vector<int> res;
    stack<TreeNode*> stk;
    TreeNode* node = root;
    if(node) stk.push(node);

    while (!stk.empty()) {
        node = stk.top();
        stk.pop();
        if (node != nullptr) {
            stk.push(node);                         // 中
            stk.push(nullptr);
            if (node->right) stk.push(node->right); // 右
            if (node->left) stk.push(node->left);   // 左
        } else {
            node = stk.top();
            stk.pop();
            res.push_back(node->val);
        }
    }

    return res;
}

总结：二叉树的递归写法中，“访问当前节点”操作和“取值”操作不是无缝衔接的，所以每个节点要入栈两次，第二次取出时再取值。为了标记两次入栈的不同，第二次入栈时，紧接着压入一个 null 指针，这样取出一个节点时，就能够知道现在应该将左右子树入栈，还是取出当前节点的值。