算法数据结构

树的遍历：整合子树的复杂信息

字数统计: 1.3k字   |   阅读时长: 4分
, ,

摘要: 树形前缀就是祖先链

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

各位好，在对树进行遍历时，假设当前的节点为 $u$，如果要在 $u$ 这个节点上进行一些复杂的计算，经常需要两个方面的信息：

  • 向上走，是 $u$ 的祖先链，我们要在 $u$ 节点做的计算依赖祖先链的信息。
  • 向下走，是以 $u$ 作为根的子树，我们要在 $u$ 节点做的计算依赖各个子树的信息。

此前我们处理过很多祖先链信息的问题，处理过程与前序遍历比较像：

本文我们来处理一个整合子树信息的问题，处理过程与后序遍历比较像，下面之前处理过的一个类似问题，可以参考：

关于树的前中后序遍历的基础知识，可以参考文章 树的前序、中序、后序遍历与DFS搜索。

题目

给你一棵二叉树的根节点 root ，二叉树中节点的值 互不相同 。另给你一个整数 start 。在第 0 分钟，感染 将会从值为 start 的节点开始爆发。

每分钟，如果节点满足以下全部条件，就会被感染：

  • 节点此前还没有感染。
  • 节点与一个已感染节点相邻。

返回感染整棵树需要的分钟数。

提示：

1
2
3
4
树中节点的数目在范围 [1, 1e5] 内
1 <= Node.val <= 1e5
每个节点的值 互不相同
树中必定存在值为 start 的节点

示例 1：

输入：root = [1,5,3,null,4,10,6,9,2], start = 3
输出：4
解释：节点按以下过程被感染：

  • 第 0 分钟：节点 3
  • 第 1 分钟：节点 1、10、6
  • 第 2 分钟：节点5
  • 第 3 分钟：节点 4
  • 第 4 分钟：节点 9 和 2
    感染整棵树需要 4 分钟，所以返回 4 。

示例 2：
输入：root = [1], start = 1
输出：0
解释：第 0 分钟，树中唯一一个节点处于感染状态，返回 0 。

题解

算法: 树的遍历

从 start 可以向两个方向传播，一个是向下，也就是向以 start 为根的子树去传播；另一个是向上，也就是 start 的祖先链上去传播。

对于向子树传播这一支，只需要知道以 start 为根的子树的深度即可。对于向祖先链传播的这一支，就比较复杂了，下面我们详细分析。

考虑 start 到 root 回溯的过程，假设当前回溯到 u 节点，此时 u 节点的左子树 u.left 和右子树 u.right 都已经处理完，start 可能在左子树，也可能在右子树：

  • 如果 start 在左子树，那么从 u 到 start 的距离加上 u 的右子树的深度就是一个可能的候选答案。
  • 如果 start 在右子树，那么从 u 到 start 的距离加上 u 的左子树的深度就是一个可能的候选答案。

对于 start 到 root 祖先链上的每个点 u，都有这样的一个候选答案。

因此大致的思路就是如果找到了 start，则 start 子树的深度是一个可能的答案，在与 start 到 root 回溯的过程中的每个候选答案共同取最大值即可。

综上，DFS 的 u 节点在回溯时要返回的信息如下：

  • u 所代表的子树中是否存在 start，若存在则为 True；
  • 若 u 子树中存在 start，返回 u 到 start 的距离，若 u 子树中不存在 start，则返回 u 子树的深度。

DFS 中对当前节点 u 的处理过程如下：

  • step1：先递归地得到 u 的子树的结果 s_left, h_left 和 s_right, h_right；
  • step2：如果 u 的值为 start，则初始化 ans 的值为 u 子树的深度。返回 (True, 1)；
  • step3：如果 start 在 u 的左子树，返回 True, h_left + 1；如果 start 在 u 的右子树，返回 True, h_right + 1；如果 start 不在 u 子树中，返回 False, max(h_left, h_right) + 1。

代码 (Python)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class Solution:
    def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:

        def dfs(node: Optional[TreeNode]) -> int:
            # 整合子树的信息：
            # 1. 子树 u 是否含 start
            # 2. 子树 u 的深度，若子树含 start，则记 u 到 start 的距离
            h_left = h_right = 0
            s_left = s_right = False
            if node.left:
                s_left, h_left = dfs(node.left)
            if node.right:
                s_right, h_right = dfs(node.right)

            nonlocal ans
            if node.val == start:
                # 找到 start 初始化
                ans = max(h_left, h_right)
                return True, 1

            if s_left or s_right:
                ans = max(ans, h_left + h_right)

            s = s_left or s_right
            h = max(h_left, h_right) + 1
            if s_left:
                h = h_left + 1
            if s_right:
                h = h_right + 1
            return s, h

        ans = 0
        dfs(root)

        return ans
Share