教你如何迭代地遍历二叉树

为何要迭代？

二叉树遍历是一个非常常见的操作，无论是中序遍历、先序遍历还是后续遍历，都可以用递归的方法很好地完成，但是相对来说迭代的方法难度就高不少，而且除此之外，迭代地遍历树至少有两个现实意义的优点：

1.比递归节省空间，递归是用栈实现的，因此如果树的高度h很大的话，递归很有可能会造成栈溢出

2.迭代的代码利用循环，而循环可以用循环不变量来证明代码的正确性

我们现在就分别详解这几个迭代地遍历树的方法。

 

先序遍历

我们先来看看先序遍历，因为在这个里面先序遍历用迭代实现是最简单的。我们先看递归的代码：

class Solution:
    # @param root, a tree node
    # @return a list of integers

    def preorder(self,root):
        ans=[]
        self.dfs(root, ans)
        return ans
    
    def dfs(self,root,ans):
        if not root:
            return
        ans.append(root.val)
        self.dfs(root.left, ans)
        self.dfs(root.right,ans)

先序遍历就是先处理根节点，再处理左右子女节点，因此用迭代实现时，我们只要处理完根节点之后，把左右子女按照先右子女、后左子女的顺序推入栈来保证这个处理顺序就可以了。这个代码也很好理解：

class Solution:
    # @param root, a tree node
    # @return a list of integers

    def preorderTraversal(self, root):
        ans=[]
        stack=[]
        if root==None:
            return ans
        stack.append(root)
        while len(stack)>0:
            top=stack.pop()
            ans.append(top.val)
            if top.right:
                stack.append(top.right)
            if top.left:
                stack.append(top.left)
        return ans

 

中序遍历

中序遍历和后续遍历递归的代码与先序遍历没有什么不同，这里就不重复了，我们直接思考迭代的方案。我们知道中序遍历是按照左、中、右的顺序遍历树，那么我们就要先找到“还没访问的节点中最左的节点L”，再找到L的右子女R，再继续按照这个步骤处理R。如何找到L？我们可以用一个变量cur来指示接下来要处理的节点，如果cur不是Null，我们就把cur压入栈，并把cur设为cur.left，反之，我们就知道现在栈顶的是L。

想清楚这些之后我们就可以开始写代码了：

class Solution:
    # @param root, a tree node
    # @return a list of integers
    def inorder(self,root):
        ans=[]
        if not root:
            return ans
        stack=[]
        cur=root
        done=False
        # 循环不变量: cur: 还没有被处理的最左子树的根节点, stack: 所有已被访问，但是因为顺序不满足中序遍历而还没有处理的节点。　　　　　　
        # 循环维持条件: stack不为空或者cur不是Null，根据cur和stack的定义，这个条件表明还有没有被处理的节点。
        # 初始化: cur是root节点。是整个树的根节点，符合cur定义。stack是空的，因为还没有访问任何节点。
        # 循环结束结果: cur是Null，有两种情况，一种是栈顶现在是最左未访问子树根节点，另一种是cur已经到了最右节点的左子女。结合stack是空的条件，说明是第二种情况，整个树已经完成了中序
        # 遍历
        while len(stack)>0 or cur:
            if cur:
                stack.append(cur)
                cur=cur.left
            else:
                cur=stack.pop()
                ans.append(cur.val)# cur是最左可以节点，要弹出栈并且遍历。保持了stack的特性。
                cur=cur.right        
        return ans
            

循环里判断cur是不是Null，如果不是Null，我们把cur压入栈，继续向左寻找“可能的更左的节点”。如果cur是Null，说明现在栈顶是未被访问的最左节点L，这时候L是下一个应该被遍历的节点，因此我们把L弹出栈，并把L加进遍历结果数组ans，再把cur设为cur.right，继续处理L的右子树的节点（想象一下，此时L的右子树的节点是除了L之外未被访问的最左节点，这时候把cur设为cur.right就相当于把从右子树根节点开始处理右子树的节点）。程序注释里包含了程序的正确性证明，我们也可以从直觉的角度做如下解释：每个节点的左子女都后被压入栈，因此每个节点的左子女都先出栈，先被处理；另外每个节点都在自己被处理并且弹出栈之后才开始把右子女压入栈，因此从直觉上来说，每个节点都保证了左->中->右的处理顺序。

 

后序遍历

后序遍历的迭代在这三个里面是最难的一个，有一种比较通用的方法（也可以用在前面两种遍历上），试想一下我们在做这几种二叉树遍历的时候最困难的问题在于我们到达一个节点的时候，我们不知道我们正在从上向下遍历还是从下往上遍历，比如说，在中序遍历，如果我们正在从上往下遍历，那我们应该先处理子女节点再处理根节点；如果我们正在从下往上遍历，那说明我们已经处理完了子女节点，可以处理根节点了。因此我们引入变量pre表示我们之前处理的节点，top表示我们正在处理的节点；我们维持pre和top是父亲子女关系或者子女父亲关心，通过判断他们是哪种关系我们就知道遍历的方向，继而能做出相应的操作：

    def postorderTraversal(self, root):
        ans=[]
        if not root:
            return ans
        pre=None
        stack=[]
        stack.append(root)
        while len(stack)>0:
            top=stack[len(stack)-1]
            if not pre or pre.left==top or pre.right==top:
                if top.left:
                    stack.append(top.left)
                elif top.right:
                    stack.append(top.right)
                else:
                    ans.append(top.val)
                    stack.pop()
            elif top.left==pre:
                if top.right:
                    stack.append(top.right)
                else: 
                    ans.append(top.val)
                    stack.pop()
            elif top.right==pre:
                ans.append(top.val)
                stack.pop()
            pre=top
        return ans

循环中分为三种情况：

1.pre是top的父亲节点，说明正在从上往下遍历，这种情况下如果top.left不是空的，那我们应该先处理左子女，因此把top.left压入栈，为了维持栈顶元素和pre的父亲子女关系，我们暂时不处理右子女，留到从下往上遍历时再处理。但是如果top.left是空而且top.right不为空，我们就把top.right压入栈，如果子女都为空的话，我们就可以把当前节点弹出并放入结果数组

2.pre是top的左子女，说明正在从下往上遍历且左子树已处理完成，这时如果top的右子女不为空，我们就把top.right压入栈，继续处理右子女，否则我们就弹出当前节点并放入结果数组

3.pre是top的右子女，说明正在从下往上遍历且左子树、右子树均处理完成（每次都是先处理左子树，因此右子树处理完成说明左子树也已处理完成），这时我们就弹出当前节点并放入结果数组。

我们可以证明，循环中所有条件下的操作维持了如下一个不变量：【pre和cur是父亲子女或者子女父亲关系，stack保存了所有未被处理的节点中的最左路径】

后序遍历中用到的这种方法同样适用于先序遍历和中序遍历，具体的代码留给读者完成。

最后这里再提供一种比较讨巧的后续遍历的方法，双栈法，用两个栈来处理：

    def postorderTraversal2(self, root):    
        ans=[]
        if not root:
            return ans
        s1=[]
        s2=[]
        s1.append(root)
        # {inv P: s1 subtrees and s2 nodes contains all nodes in the original tree. and s1's nodes is poped in anti-postorder }
        # {initialization: s1 contains root subtree , s2 is empty. root node will be poped first, it is anti-postorder}
        # { !B ^ P => s1 is empty, so s2 contains all nodes and s2's nodes is in anti-postorder since s2 add nodes in exact order how nodes are poped from s1. so pop s2 will get a post order sequence}
        # { why is invariant true?: s2 add the node poped from s1 and push children of this node back to s1. so P1 is maintained. since the node's right child is pushed after left node
        # so it will be poped before left node after root node, which is exact anti-postorder }
        while len(s1)>0:
            top=s1.pop()
            s2.append(top)
            if top.left:
                s1.append(top.left)            
            if top.right:
                s1.append(top.right)
        while len(s2)>0:
            top=s2.pop()
            ans.append(top.val)
        return ans
        

参考文献

 [1] Binary Tree Post-Order Traversal Iterative Solution  http://leetcode.com/2010/10/binary-tree-post-order-traversal.html