算法专题:二叉树的遍历

二叉树的遍历想必大家都不陌生,主要有三种遍历方式:前序遍历(pre-order traversal),中序遍历(in-order ~)还有后序遍历(post-order ~)。所谓前中后,就是遍历根节点相对于遍历子树的顺序。这里假设读者已经了解了最基础的三种遍历的定义,不再赘述。
这篇文章的目的就是用代码实现二叉树的遍历。不仅要实现,还要多种方法实现,还要让读者能记住一些写法。因为二叉树遍历应用实在很广,是很基础很实用的算法,不得不熟悉。
本篇文章编程语言为Kotlin(Python:???)。不熟悉这门语言的也没关系,大致应该还是能看懂。
树节点的定义:

class TreeNode(var value: Int = 0, var left: TreeNode? = null, var right: TreeNode? = null)

简单的递归

二叉树简直就是为递归而生的。一棵二叉树可以看成左子树,根节点和右子树,也就是说往往可以对左右子树递归调用。
对于遍历,递归无疑是最简单的解决方式。递归的思维是这样的:假设递归函数的作用就是以某种方式遍历树,那么再来看左子树,根节点和右子树,无非就是按照某个顺序对左右子树调用并把根节点放到结果里。
还是看代码吧,真的很简单(x?.let{}表示如果x不为null就执行括号里面的内容,it在括号内指代x):

    // T:O(n) S:O(h)
    fun inOrderRecur(root: TreeNode): List<Int> {
        val ret = mutableListOf<Int>()
        root.left?.let { ret.addAll(inOrderRecur(it)) }
        ret.add(root.value)
        root.right?.let { ret.addAll(inOrderRecur(it)) }
        return ret
    }

    // T:O(n) S:O(h)
    fun preOrderRecur(root: TreeNode): List<Int> {
        val ret = mutableListOf<Int>()
        ret.add(root.value)
        root.left?.let { ret.addAll(preOrderRecur(it)) }
        root.right?.let { ret.addAll(preOrderRecur(it)) }
        return ret
    }

    // T:O(n) S:O(h)
    fun postOrderRecur(root: TreeNode): List<Int> {
        val ret = mutableListOf<Int>()
        root.left?.let { ret.addAll(postOrderRecur(it)) }
        root.right?.let { ret.addAll(postOrderRecur(it)) }
        ret.add(root.value)
        return ret
    }

可以看到,递归代码真的很短,也非常好记忆。
不过递归毕竟是递归,这里也并没有tail recursion,效率总是不如迭代方法的。如何用非递归方法来实现遍历呢?

模仿递归的栈解法

一个很科学的思路就是把递归“翻译”成迭代。既然我已经有了这么简单方便的递归方法,我只需要把它们转化成迭代方法,不就可以了吗?
那么这里的关键就是:如何“翻译”?
要模仿递归的思路,首先需要栈,因为递归虽然没有声明栈,但是系统提供了栈。
再仔细思考一下会发现节点的遍历顺序很重要,有些节点会留在之后再遍历。
有没有一种类似于递归的栈解法,用很相似的代码来实现三种遍历?
有的。诀窍就是引入一个表示是否已经遍历过的布尔值isVisited。递归其实也是有这么一个过程的,从第一次调用,到子问题,然后再回来。
节点第一次进栈时,布尔值为false,之后再出现的时候,假设是true,那么添加到结果里,否则改为true再按照遍历顺序进栈。
代码如下:

    // T:O(n) S:O(h)
    fun inOrder(root: TreeNode): List<Int> {
        val stk = Stack<Pair<TreeNode?, Boolean>>()
        val ret = mutableListOf<Int>()
        stk.push(Pair(root, false))
        while (stk.isNotEmpty()) {
            val p = stk.pop()
            val cur = p.first
            val isVisited = p.second
            if (cur == null) {
                continue
            } else {
                if (isVisited) {
                    ret.add(cur.value)
                } else {
                    stk.push(Pair(cur.right, false))
                    stk.push(Pair(cur, true))
                    stk.push(Pair(cur.left, false))
                }
            }
        }
        return ret
    }

    // T:O(n) S:O(h)
    fun preOrder(root: TreeNode): List<Int> {
        val stk = Stack<Pair<TreeNode?, Boolean>>()
        val ret = mutableListOf<Int>()
        stk.push(Pair(root, false))
        while (stk.isNotEmpty()) {
            val p = stk.pop()
            val cur = p.first
            val isVisited = p.second
            if (cur == null) {
                continue
            } else {
                if (isVisited) {
                    ret.add(cur.value)
                } else {
                    stk.push(Pair(cur.right, false))
                    stk.push(Pair(cur.left, false))
                    stk.push(Pair(cur, true))
                }
            }
        }
        return ret
    }

    // T:O(n) S:O(h)
    fun postOrder(root: TreeNode): List<Int> {
        val stk = Stack<Pair<TreeNode?, Boolean>>()
        val ret = mutableListOf<Int>()
        stk.push(Pair(root, false))
        while (stk.isNotEmpty()) {
            val p = stk.pop()
            val cur = p.first
            val isVisited = p.second
            if (cur == null) {
                continue
            } else {
                if (isVisited) {
                    ret.add(cur.value)
                } else {
                    stk.push(Pair(cur, true))
                    stk.push(Pair(cur.right, false))
                    stk.push(Pair(cur.left, false))
                }
            }
        }
        return ret
    }

这种解法代码稍长一点,好在三种解法几乎一模一样,除了进栈的顺序不同。这里进栈的顺序就是对遍历顺序倒着来,因为先进后出。

各显神通的栈解法

另一种思路是抛弃递归的思路,而分别针对三种递归的特性来用栈实现遍历。这样一来,每种解法都显得更加优雅,但是不同解法之间又不太一样,至少像上面那样只需要记一种解法是不可能的了。
代码:

    // T:O(n) S:O(h)
    fun inOrderWithStack(root: TreeNode): List<Int> {
        val stk = Stack<TreeNode>()
        val ret = mutableListOf<Int>()
        var cur: TreeNode? = root
        while (cur != null || stk.isNotEmpty()) {
            while (cur != null) {
                stk.push(cur)
                cur = cur.left
            }
            cur = stk.pop()
            ret.add(cur.value)
            cur = cur.right
        }
        return ret
    }

    // T:O(n) S:O(h)
    fun preOrderWithStack(root: TreeNode): List<Int> {
        val stk = Stack<TreeNode>()
        val ret = mutableListOf<Int>()
        stk.push(root)
        while (stk.isNotEmpty()) {
            val cur = stk.pop()
            ret.add(cur.value)
            cur.right?.let { stk.push(it) }
            cur.left?.let { stk.push(it) }
        }
        return ret
    }

    // T:O(n) S:O(h)
    fun postOrderWithDoubleStacks(root: TreeNode): List<Int> {
        val stk1 = Stack<TreeNode>()
        val stk2 = Stack<TreeNode>()
        stk1.push(root)
        while (stk1.isNotEmpty()) {
            val cur = stk1.pop()
            stk2.push(cur)
            cur.left?.let { stk1.push(it) }
            cur.right?.let { stk1.push(it) }
        }
        return stk2.map { it.value }.reversed()
    }

三种解法虽然不同,但又有千丝万缕的联系。中序、前序都只需要一个栈,后序需要两个栈,但后序和前序其实很相似。

挑战:Morris遍历

说到树的遍历,不提及Morris遍历总是不完整的。
Morris遍历能以O(n)的时间复杂度和O(1)的空间复杂度来实现三种遍历,相对应的代价就是代码更为复杂。
在实际面试中,除非明确要求,或者把握很大,且又是瓶颈,不然没必要为了这O(n)的空间复杂度来用Morris,尤其是最复杂的后序遍历。当然,做出解答之后也许可以提一下,说用Morris能得到更好的空间复杂度。
既然是O(1)的空间复杂度,说明Morris算法不用栈,那么如何回到以前遍历过的节点呢?
答案就是指针。
原本树的叶子节点,左右两个指针都是空的,Morris就在这方面“耍花招”,让其右指针指向遍历顺序该返回的节点。

中序

以中序遍历为例,考虑以下这棵树:
    1
    /
   3
  / 
 4  2
 /
5
遍历的结果应该是54321,也就是说5的返回点是4,4的返回点是3,2的返回点应该是1。
还是来看代码吧:

    // T:O(n) S:O(1)
    fun inOrderMorris(root: TreeNode): List<Int> {
        val ret = mutableListOf<Int>()
        var cur: TreeNode? = root
        while (cur != null) {
            if (cur.left == null) {
                ret.add(cur.value)
                cur = cur.right
            } else {
                var node: TreeNode = cur.left!!
                while (node.right != null && node.right != cur)
                    node = node.right!!
                if (node.right == null) {
                    node.right = cur
                    cur = cur.left
                } else {
                    ret.add(cur.value)
                    node.right = null
                    cur = cur.right
                }
            }
        }
        return ret
    }

解释一下这个代码对于上面那棵树的执行过程:

  • cur为1,此时把2的右指针指向1,cur往左走一步;
  • cur为3,把4的右指针指向3,cur往左走一步;
  • cur为4,把5的右指针指向4,cur往左走一步;
  • cur为5,此时把5加入结果,通过右指针返回4;
  • cur为4,此时5的右指针还指向4,说明4是第二次来了,把5的右指针复原,把4加入结果,cur往右走一步;
  • cur为3,类似地把4的右指针复原,把3加入结果,cur往右走一步;
  • cur为2,把2加入结果,通过右指针返回1;
  • cur为1,把2的右指针复原,把1加入结果,往右走一步;
  • cur为null,结束

可以发现,Morris遍历其实要走两遍节点,第一遍主要是架设“后路”。
仔细思考,根节点左子树的最右节点,不就是中序遍历里面根节点的前一个元素吗?因为是左子树的最右节点,说明左子树的中序已经完成,接下来就是返回根节点了。
第二遍就是获得结果并复原指针。通过是否有“后路”来确定是不是第二遍,然后再做行动。
总而言之,非常巧妙。
Morris本人其实就提出了中序遍历,其他人根据中序遍历,推出了前序和后序遍历的Morris方法。

前序

再来看看前序遍历:

    // T:O(n) S:O(1)
    fun preOrderMorris(root: TreeNode): List<Int> {
        val ret = mutableListOf<Int>()
        var cur: TreeNode? = root
        while (cur != null) {
            if (cur.left == null) {
                ret.add(cur.value)
                cur = cur.right
            } else {
                var node = cur.left
                while (node?.right != null && node.right != cur)
                    node = node.right
                if (node?.right == null) {
                    ret.add(cur.value)
                    node?.right = cur
                    cur = cur.left
                } else {
                    node.right = null
                    cur = cur.right
                }
            }
        }
        return ret
    }

这个代码明显和中序非常像。总的来说,架设“后路”的思路还是一样的,只不过加入结果的时机不同。
还是以上面那棵树为例:

  • cur为1,先把1加入结果,再让2指向1,cur往左走一步;
  • cur为3,把3加入结果,让4指向3,cur往左走一步;
  • cur为4,把4加入结果,让5指向4,cur往左走一步;
  • cur为5,把5加入结果,往右走一步,通过后路返回4;
  • cur为4,把5的后路复原,往右走一步,通过后路返回3;
  • cur为3,把4的后路复原,往右走一步;
  • cur为2,把2加入结果,往右走一步,通过后路返回1;
  • cur为1,把2的后路复原,往右走一步;
  • cur为null,结束

可以看到,中序和前序的行动轨迹是完全一样的,而且都在最左端加入结果。区别在于,中序是回来的时候加入结果,前序是第一次就加入结果。

后序

Morris的后序又是另一个故事了,比前序和中序都要复杂很多。

    // T:O(n) S:O(1)
    fun postOrderMorris(root: TreeNode): List<Int> {
        val dummy = TreeNode()
        dummy.left = root
        val ret = mutableListOf<Int>()
        var cur: TreeNode? = dummy
        while (cur != null) {
            if (cur.left == null) {
                cur = cur.right
            } else {
                var node = cur.left
                while (node?.right != null && node.right != cur)
                    node = node.right
                if (node?.right == null) {
                    node?.right = cur
                    cur = cur.left
                } else {
                    ret.addAll(traceBack(cur.left!!, node))
                    node.right = null
                    cur = cur.right
                }
            }
        }
        return ret
    }

    fun traceBack(from: TreeNode, to: TreeNode): List<Int> {
        val ret = mutableListOf<Int>()
        var cur: TreeNode = from
        while (cur !== to) {
            ret.add(cur.value)
            cur = cur.right!!
        }
        ret.add(to.value)
        return ret.reversed()
    }

试推演执行步骤如下:

  • cur为dummy,将1的右指针指向dummy,cur往左走一步;
  • cur为1,将2的右指针指向1,cur往左走一步;
  • cur为3,将4的右指针指向3,cur往左走一步;
  • cur为4,将5的右指针指向4,cur往左走一步;
  • cur为5,此时直接通过后路返回4;
  • cur为4,将5加入结果,并复原5的右指针,通过后路返回3;
  • cur为3,类似地将4加入结果,复原4的右指针,向右走一步;
  • cur为2,通过后路返回1;
  • cur为1,将2、3加入结果,复原2的右指针,向右走一步;
  • cur为dummy,将1加入结果,复原1的右指针,向右走一步;
  • cur为null,结束

架设路径仍旧是老套路,只是这次不仅使用了一个dummy node,而且还需要一个子函数来配合获取结果,确实复杂了很多。

推荐阅读更多精彩内容