二叉树主题

2023-08-06

前/中/后序遍历不是三个算法,是同一个递归框架里”打印根”那一行放在哪。递归一行搞定;迭代用栈,是在手动模拟系统调用栈。BFS 换成队列就是层序。记住这个,二叉树遍历题不会再混。

直觉 — 前/中/后序的差别只是根的位置

写出标准的二叉树递归骨架:

1
2
3
4
5
6
7
def dfs(node):
    if node is None: return
    # (A) 这里访问 → 前序
    dfs(node.left)
    # (B) 这里访问 → 中序
    dfs(node.right)
    # (C) 这里访问 → 后序

根的访问语句夹在两次递归调用 (左, 右) 之间的哪个位置,就决定了遍历名字:

遍历 顺序 根的位置 适合干什么
前序 根 → 左 → 右 两次递归之前 自顶向下传参数 (路径/前缀)
中序 左 → 根 → 右 两次递归之间 BST 升序输出
后序 左 → 右 → 根 两次递归之后 自底向上聚合 (高度/直径)

层序遍历是另一类: 不再递归,改用队列,一层一层往外铺。本质上就是 BFS。

动画 — 同一棵 7 节点树,三种顺序

下面这棵完全二叉树,根是 1,叶子是 4,5,6,7。点动画上的按钮切换遍历方式,观察访问顺序: 前序从根开始往左下扎,中序从最左下角开始,后序最后一个才到根。一眼看出”根的位置”差别。

关键观察: 前序第一个一定是根,后序最后一个一定是根;中序里根的位置取决于左子树大小。这就是”前序+中序唯一确定一棵树”的根基。

四种遍历 — 递归 + 迭代

递归 — 一行之差

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def preorder(node):
    if node is None: return
    print(node.val)            # 根
    preorder(node.left)        # 左
    preorder(node.right)       # 右

def inorder(node):
    if node is None: return
    inorder(node.left)
    print(node.val)            # 根在中间
    inorder(node.right)

def postorder(node):
    if node is None: return
    postorder(node.left)
    postorder(node.right)
    print(node.val)            # 根在最后

三个函数只有 print 这一行的位置不同。这就是全部秘密。

迭代 — 用栈模拟系统调用栈

为什么需要栈? 递归之所以能”先去左子树,回来再处理右子树”,是靠系统调用栈帮你把”当前位置”保存了。改写成迭代,就必须自己用一个显式栈,把”下一步要回来做什么”压进去。

前序最简单 (访问完根就立刻能出栈):

1
2
3
4
5
6
7
8
9
def preorder_iter(root):
    if root is None: return
    stack = [root]
    while stack:
        node = stack.pop()
        print(node.val)
        # 右先入栈,左后入栈 → 左先出栈,符合"根左右"
        if node.right: stack.append(node.right)
        if node.left:  stack.append(node.left)

中序难一点 — 必须先把左链全部压进去,再回头处理:

1
2
3
4
5
6
7
8
9
def inorder_iter(root):
    stack, node = [], root
    while stack or node:
        while node:                  # 一路向左压栈
            stack.append(node)
            node = node.left
        node = stack.pop()           # 弹出 = 该访问"根"了
        print(node.val)
        node = node.right            # 接着去右子树

后序最绕。最简洁的写法: 按”根右左”的顺序前序遍历,结果反转,得到”左右根”:

1
2
3
4
5
6
7
8
9
def postorder_iter(root):
    if root is None: return []
    stack, out = [root], []
    while stack:
        node = stack.pop()
        out.append(node.val)
        if node.left:  stack.append(node.left)   # 注意顺序反了
        if node.right: stack.append(node.right)
    return out[::-1]                              # 反转 → 左右根

层序 = BFS,用队列

把栈换成队列,深度优先变广度优先:

1
2
3
4
5
6
7
8
9
10
11
12
13
from collections import deque

def level_order(root):
    if root is None: return
    q = deque([root])
    while q:
        n = len(q)                   # 锁定本层长度,才能按层切分
        for _ in range(n):
            node = q.popleft()
            print(node.val, end=' ')
            if node.left:  q.append(node.left)
            if node.right: q.append(node.right)
        print()                       # 本层结束换行

n = len(q) 这一句是层序的灵魂 — 没有它就是普通 BFS,有了它才能”一层一层”。

应用 — 三个必知

1. BST 中序 = 升序

BST 定义: 左子树全 < 根 < 右子树全。中序”左根右”恰好把这个偏序展开成全序。验证 BST、第 k 小、范围求和,全靠这一条。

1
2
3
4
5
6
7
8
9
def is_valid_bst(root):
    prev = [-float('inf')]
    def dfs(node):
        if node is None: return True
        if not dfs(node.left): return False
        if node.val <= prev[0]: return False     # 中序应严格递增
        prev[0] = node.val
        return dfs(node.right)
    return dfs(root)

2. 前序 + 中序 → 唯一构造

前序的第一个是根,拿这个根去中序里切一刀: 左边是左子树的中序,右边是右子树的中序。递归下去就建好了。

1
2
3
4
5
6
7
8
9
10
11
def buildTree(preorder, inorder):
    idx = {v: i for i, v in enumerate(inorder)}  # O(1) 定位
    def build(pl, pr, il, ir):
        if pl > pr: return None
        root = TreeNode(preorder[pl])
        m = idx[preorder[pl]]
        left_size = m - il
        root.left  = build(pl+1, pl+left_size, il, m-1)
        root.right = build(pl+left_size+1, pr, m+1, ir)
        return root
    return build(0, len(preorder)-1, 0, len(inorder)-1)

注意: 只有前序、只有中序、或只有后序,都不能唯一确定一棵树。必须有中序 + (前序或后序) 这一对。

3. Morris 遍历 — O(1) 空间

递归 O(h) 栈空间,迭代 O(h) 栈空间。Morris 把额外空间降到 O(1) — 代价是临时改树。

核心 idea: 当前节点 cur 有左子树时,找到左子树的最右节点 (中序前驱),把它的 right 指针指向 cur (“线索”),这样从左子树最深处可以”飞”回 cur。访问完 cur 后再把线索摘掉,恢复原树。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def morris_inorder(root):
    cur = root
    while cur:
        if cur.left is None:
            print(cur.val)            # 没左子树,访问自己,转右
            cur = cur.right
        else:
            pre = cur.left            # 找左子树最右节点
            while pre.right and pre.right is not cur:
                pre = pre.right
            if pre.right is None:     # 第一次到 cur,建线索
                pre.right = cur
                cur = cur.left
            else:                     # 第二次到 cur,拆线索
                pre.right = None
                print(cur.val)
                cur = cur.right

面试题里不常考代码,但常考”如何 O(1) 空间遍历?” — 答 Morris。

经典题

LC 题目 用哪种遍历
104 二叉树的最大深度 后序 (拿子树深度)
226 翻转二叉树 前序 / 后序都行
98 验证 BST 中序 (检查升序)
543 二叉树的直径 后序 (左深+右深)
124 最大路径和 后序 + “经过我的 vs 给上面的”双值
236 最近公共祖先 后序 (左右各找一遍再判断根)
105 前序+中序建树 前序定根,中序切左右
102 层序遍历 BFS + len(q) 分层
297 序列化/反序列化 前序 (含 null) 最自然

判断用哪种的口诀: 当前节点的答案需要子树返回值 → 后序;需要从上往下传 → 前序;BST 性质 → 中序;按层处理 → 层序。

记忆法

  • 三种遍历差别只有一句话: 打印根的那行,放在 dfs(left) 和 dfs(right) 这两个调用的前/中/后。
  • 递归一行搞定,迭代要栈: 栈是在手动模拟系统调用栈,把”接下来要做什么”显式保存。
  • 栈 ↔ DFS, 队列 ↔ BFS: 数据结构一换,深度变广度。层序的 len(q) 才是”按层”。
  • 前序+中序 = 唯一一棵树: 前序定根,中序切左右,递归。
  • BST 中序 = 升序: 凡是涉及 BST 的有序性,先想中序。
  • O(1) 空间遍历 = Morris: 用叶子的空右指针临时建线索,访问完拆掉。

一句话总览: 二叉树遍历 = 递归骨架 + 一行根位置;迭代是把递归栈摊开来写;层序是把栈换成队列。