移除子树后的二叉树高度

标签: 树 深度优先搜索 广度优先搜索 数组 二叉树

难度: Hard

给你一棵 二叉树 的根节点 root ，树中有 n 个节点。每个节点都可以被分配一个从 1 到 n 且互不相同的值。另给你一个长度为 m 的数组 queries 。

你必须在树上执行 m 个 独立 的查询，其中第 i 个查询你需要执行以下操作：

  • 从树中 移除 以 queries[i] 的值作为根节点的子树。题目所用测试用例保证 queries[i] 不 等于根节点的值。

返回一个长度为 m 的数组 answer ，其中 answer[i] 是执行第 i 个查询后树的高度。

注意：

  • 查询之间是独立的，所以在每个查询执行后，树会回到其 初始 状态。
  • 树的高度是从根到树中某个节点的 最长简单路径中的边数 。

示例 1：

输入：root = [1,3,4,2,null,6,5,null,null,null,null,null,7], queries = [4]
输出：[2]
解释：上图展示了从树中移除以 4 为根节点的子树。
树的高度是 2（路径为 1 -> 3 -> 2）。

示例 2：

输入：root = [5,8,9,2,1,3,7,4,6], queries = [3,2,4,8]
输出：[3,2,3,2]
解释：执行下述查询：
- 移除以 3 为根节点的子树。树的高度变为 3（路径为 5 -> 8 -> 2 -> 4）。
- 移除以 2 为根节点的子树。树的高度变为 2（路径为 5 -> 8 -> 1）。
- 移除以 4 为根节点的子树。树的高度变为 3（路径为 5 -> 8 -> 2 -> 6）。
- 移除以 8 为根节点的子树。树的高度变为 2（路径为 5 -> 9 -> 3）。

提示：

  • 树中节点的数目是 n
  • 2 <= n <= 105
  • 1 <= Node.val <= n
  • 树中的所有值 互不相同
  • m == queries.length
  • 1 <= m <= min(n, 104)
  • 1 <= queries[i] <= n
  • queries[i] != root.val

Submission

运行时间: 369 ms

内存: 52.2 MB

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def calDepth(self, root):
        root.depth = 0
        if root.left:
            left = self.calDepth(root.left)
        else:
            left = -1
        if root.right:
            right = self.calDepth(root.right)
        else:
            right = -1
        root.depth = max(left + 1, right + 1)
        return root.depth

    def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
        ans = defaultdict(int)

        stack = [(root, 0)]
        x = 0
        while stack:
            node, h = stack.pop()
            ans[node.val] = x
            if h > x:
                x = h
            if node.right:
                stack.append((node.right, h + 1))
            if node.left:
                stack.append((node.left, h + 1))

        stack = [(root, 0)]
        x = 0
        while stack:
            node, h = stack.pop()
            if x > ans[node.val]:
                ans[node.val] = x
            if h > x:
                x = h
            if node.left:
                stack.append((node.left, h + 1))
            if node.right:
                stack.append((node.right, h + 1))

        return [ans[i] for i in queries]

Explain

题解中首先定义了一个辅助函数 calDepth 来计算每个节点的深度，并存储在节点属性 depth 中。这个函数递归地计算左右子树的深度，然后返回当前节点的深度，即左右子树深度的最大值加一。在 treeQueries 函数中，使用栈来进行深度优先搜索 (DFS)，计算每个节点作为树根时树的最大深度。这种方法通过两次遍历，第一次确定每个节点的深度，第二次使用深度优先搜索的迭代方式来计算每个节点的高度。最终，对于每个查询返回删除该节点后树的高度。

时间复杂度: O(n)

空间复杂度: O(n)

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def calDepth(self, root):
        # 计算节点的深度
        root.depth = 0
        if root.left:
            left = self.calDepth(root.left)
        else:
            left = -1
        if root.right:
            right = self.calDepth(root.right)
        else:
            right = -1
        root.depth = max(left + 1, right + 1)
        return root.depth

    def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
        ans = defaultdict(int)

        # 使用栈进行DFS，计算每个节点作为根时树的高度
        stack = [(root, 0)]
        x = 0
        while stack:
            node, h = stack.pop()
            ans[node.val] = x
            if h > x:
                x = h
            if node.right:
                stack.append((node.right, h + 1))
            if node.left:
                stack.append((node.left, h + 1))

        # 重置并再次进行DFS
        stack = [(root, 0)]
        x = 0
        while stack:
            node, h = stack.pop()
            if x > ans[node.val]:
                ans[node.val] = x
            if h > x:
                x = h
            if node.left:
                stack.append((node.left, h + 1))
            if node.right:
                stack.append((node.right, h + 1))

        # 返回每个查询的结果
        return [ans[i] for i in queries]

Explore

在`calDepth`方法中，将不存在的子树的深度设为-1是为了正确计算存在的子树的深度。如果一个节点不存在任何子树，其深度应为0（仅该节点自身）。设左右子树的深度为-1可以确保，当计算该节点的深度时，通过`max(left + 1, right + 1)`得出的结果为0。如果设置为0，在存在子树的情况下计算结果会错误地增加，从而使得深度计算偏大。

在题解给出的代码中，并没有直接重置树的状态的步骤。这里的`树回到初始状态`应该理解为每次DFS遍历使用的是树的原始结构，没有修改树本身的结构或节点。每次DFS开始都是从根节点开始，使用栈来保存遍历的状态，因此每次遍历都是独立的，不影响树的原始结构。

在`treeQueries`函数中，两次DFS遍历的目的是为了分别从不同的方向（左和右）计算每个节点在不同子树结构下的最大深度。第一次DFS遍历从左到右，计算从左子树开始的最大深度；第二次遍历从右到左，计算从右子树开始的最大深度。通过两次遍历，可以确保无论节点的子树结构如何变化，都能正确计算出每个节点作为树根时树的最大高度。

在使用栈进行DFS时，栈中元素的第二个参数`h`代表当前节点到根节点的高度（或称为深度）。这个参数有助于在遍历过程中记录从根节点到当前节点的路径长度。通过这个深度值，我们可以在每次从栈中取出节点时，判断并更新该节点或其父节点为根的子树的最大深度，从而准确计算出节点作为树根时的整体高度。