二叉树的 Morris 前序和中序遍历方法

刷 LeetCode 的时候碰到了一道题——LeetCode 94 二叉树的中序遍历,学会了一种有趣的非递归二叉树遍历方法,特此记录。

尝试使用空间复杂度 O(1) 的算法实现二叉树遍历。 虽然,递归法遍历二叉树代码简单且容易阅读和理解,但是在处理大规模数据时并不好用,会因为过多的递归调用导致栈空间不足。内存受限的情况下,这样的程序是不合要求的。我最开始解题的思路是使用 BFS 算法,或者使用辅助栈实现 DFS 算法,但是这样仍然会有 O(n) 的空间复杂度,未必就是最优解。刚好 Morris 遍历法可以解决这个问题。

要尽量减少额外内存消耗,那么常常会引入指针操作,以达成灵活的内存跳转。又因为二叉树节点除了包含数据外,一般还会保存左右子树的指针,指针的额外内存占用已经包含在输入数据构建之中。常规的递归遍历实际是为了使用额外的空间记录子树遍历路径,方便遍历操作到达子树叶子节点之后层层返回。那么,要减少额外空间的消耗,待解决的问题就是:如何使用指针操作,使得遍历操作进入子树后仍能正确返回上层节点。 这时,我们应该先模拟遍历操作的流程分析指针操作的具体步骤,不妨先暂且只分析中序遍历,重点观察遍历操作在进入子树后因何返回上层节点,以及从何种节点进行返回跳转。假设二叉树数据如下所示:

1
2
3
4
5
          1
        /   \
       2     3
      / \   /
     4   5 6

假设我们要从根节点为 2 的子树处跳转到节点 1 处,原因是我们已经完成了子树的遍历,而子树的中序遍历终止条件就是我们处理完了它的最右侧的叶子节点。那么,我们何妨如此处理中序遍历问题:

  1. 检查当前节点的左子树是否是空,如果非空则需要进行左子树遍历
  2. 在进行左子树遍历之前,先找到左子树的最右叶子节点
  3. 把步骤 2 中找到的叶子节点的右侧指针指向当前节点
  4. 进入左子树遍历

这样,例子中的二叉树会发生如下变化:

第一次叶子右指针赋值:

1
2
3
4
5
6
7
8
9
         2<----
        / \   |
       4   5  |
            \ |
             1
              \
               3
              /
             6

第二次叶子右指针赋值:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
        4<-
         \|
          2<---
           \  |
            5 |
             \|
              1
               \
                3
               /
              6

特别注意:此时,树中必然存在环路,这正是后续操作需要处理的问题,也是还原操作可以利用的条件。

当遍历操作退回上层节点时,我们又会继续检查左子树的最右叶子节点,但是因为环路的存在,实际上并不存在合法的叶子节点。但是,可以找到右子节点时当前节点的一个节点,它实际上就是我们之前进行指针赋值时操作的叶子节点。那么,只要对它进行还原操作,就能保证左子树恢复如初,具体还原操作如下:

  1. 在当前节点检查左子树是否为空,如果非空则进行还原操作
  2. 在左子树的右侧后代中查找一个节点,要求它的右子节点为步骤 1 中记录的当前节点
  3. 把步骤 2 中找到节点右指针清空

还原操作进行完毕后,按照中序遍历要求,我们处理当前节点的数据,然后进入当前节点的右子树继续进行遍历即可。

我们应该注意到,上面的算法设计思路中叶子指针跳转和还原操作的结构相同,而且节点行进路径也完全一致,可以进行合并简化。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
fun inorderTraversal(root: TreeNode) {
    curr = root
    while (curr != null) {
        if (curr.left == null) {
            process curr.val
            curr = curr.right
        } else {
            loop in curr.left to determine ->
                1. prev is the right leaf of curr.left: {  // prev.right is null
                    prev.right = curr
                    curr = curr.left
                }
                2. prev.right is curr {
                    prev.right = null
                    process curr.val
                    curr = curr.right
                }
        }
    }
}

而确定最右叶子节点的伪代码也很好构建:

1
2
3
4
5
6
...
    prev = curr.left
    while (prev.right != null and prev.right != curr) {
        prev = prev.right
    }
...

二叉树的节点代码如下:

1
2
3
4
class TreeNode(var `val`: Int) {
    var left: TreeNode? = null
    var right: TreeNode? = null
}

中序 Morris 遍历的代码如下所示:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
fun inorderTraversal(root: TreeNode?): List<Int> {
    var curr = root
    val re = ArrayList<Int>()
    while (curr != null) {
        if (curr.left == null) {
            re.add(curr.`val`)
            curr = curr.right
        } else {
            var prev = curr.left!!
            while (prev.right != null && prev.right != curr) prev = prev.right!!
            if (prev.right == null) {
                prev.right = curr
                curr = curr.left
            } else {
                prev.right = null
                re.add(curr.`val`)
                curr = curr.right
            }
        }
    }
    return re
}

前序遍历的分析不再赘述,只展示代码。 相应的前序遍历代码如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
private fun preorderTraversal(root: TreeNode?): List<Int> {
    var curr = root
    val re = ArrayList<Int>()
    while (curr != null) {
        if (curr.left == null) {
            re.add(curr.`val`)
            curr = curr.right
        } else {
            var prev = curr.left!!
            while (prev.right != null && prev.right != curr) prev = prev.right!!
            if (prev.right == null) {
                re.add(curr.`val`)
                prev.right = curr
                curr = curr.left
            } else {
                prev.right = null
                curr = curr.right
            }
        }
    }
    return re
}

注意:复原操作时无需处理当前节点的数据,只需要还原指针即可,所以处理数据的操作必须如此,不可提前到外层循环中。