教你如何迭代地遍歷二叉樹

為何要迭代？

二叉樹遍歷是一個非常常見的操作，無論是中序遍歷、先序遍歷還是后續遍歷，都可以用遞歸的方法很好地完成，但是相對來說迭代的方法難度就高不少，而且除此之外，迭代地遍歷樹至少有兩個現實意義的優點：

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