Hard

题目描述

给你一个有 n 个节点的二叉树的根节点 root，每个节点都有一个从 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 <= 10⁵
  • 1 <= Node.val <= n
  • 树中所有值都是唯一的
  • m == queries.length
  • 1 <= m <= min(n, 10⁴)
  • 1 <= queries[i] <= n
  • queries[i] != root.val

解题思路

这道题的关键在于预处理，避免对每个查询都重新计算整棵树。

思路分析：

核心思想是预计算每个节点被移除后的树高度。对于任意节点的移除，新的树高度等于从根节点出发，不经过被移除节点的最长路径。

算法步骤：

  1. 第一次DFS：计算每个节点的深度（距离根节点的距离）和每个子树的高度
  2. 第二次DFS：自顶向下计算每个节点被移除后的答案
    • 对于当前节点，考虑其父节点能提供的最大深度
    • 同时考虑兄弟子树能提供的最大深度
    • 这两者的最大值就是移除当前节点后的树高度

关键观察：

  • 当移除某个节点时，剩余树的高度来源于两部分：
    1. 从根到该节点父节点的路径继续向其他子树延伸
    2. 该节点的兄弟子树

通过两次DFS遍历，第一次自底向上计算子树信息，第二次自顶向下传递父节点信息，可以在O(n)时间内预处理出所有答案，每个查询O(1)响应。

代码实现

class Solution {
public:
    vector<int> treeQueries(TreeNode* root, vector<int>& queries) {
        unordered_map<int, int> nodeDepth, subtreeHeight, answer;
        
        // First DFS: calculate depth and subtree height
        function<int(TreeNode*, int)> dfs1 = [&](TreeNode* node, int depth) -> int {
            if (!node) return -1;
            
            nodeDepth[node->val] = depth;
            int leftHeight = dfs1(node->left, depth + 1);
            int rightHeight = dfs1(node->right, depth + 1);
            
            return subtreeHeight[node->val] = max(leftHeight, rightHeight) + 1;
        };
        
        // Second DFS: calculate answer for each node
        function<void(TreeNode*, int)> dfs2 = [&](TreeNode* node, int maxDepthFromParent) {
            if (!node) return;
            
            answer[node->val] = maxDepthFromParent;
            
            int leftHeight = node->left ? subtreeHeight[node->left->val] : -1;
            int rightHeight = node->right ? subtreeHeight[node->right->val] : -1;
            
            // For left child: max depth comes from parent or right sibling
            dfs2(node->left, max(maxDepthFromParent, nodeDepth[node->val] + 1 + rightHeight));
            // For right child: max depth comes from parent or left sibling  
            dfs2(node->right, max(maxDepthFromParent, nodeDepth[node->val] + 1 + leftHeight));
        };
        
        dfs1(root, 0);
        dfs2(root, 0);
        
        vector<int> result;
        for (int query : queries) {
            result.push_back(answer[query]);
        }
        return result;
    }
};

class Solution:
    def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
        node_depth = {}
        subtree_height = {}
        answer = {}
        
        def dfs1(node, depth):
            if not node:
                return -1
            
            node_depth[node.val] = depth
            left_height = dfs1(node.left, depth + 1)
            right_height = dfs1(node.right, depth + 1)
            
            subtree_height[node.val] = max(left_height, right_height) + 1
            return subtree_height[node.val]
        
        def dfs2(node, max_depth_from_parent):
            if not node:
                return
            
            answer[node.val] = max_depth_from_parent
            
            left_height = subtree_height.get(node.left.val, -1) if node.left else -1
            right_height = subtree_height.get(node.right.val, -1) if node.right else -1
            
            dfs2(node.left, max(max_depth_from_parent, node_depth[node.val] + 1 + right_height))
            dfs2(node.right, max(max_depth_from_parent, node_depth[node.val] + 1 + left_height))
        
        dfs1(root, 0)
        dfs2(root, 0)
        
        return [answer[query] for query in queries]

public class Solution {
    private Dictionary<int, int> nodeDepth = new Dictionary<int, int>();
    private Dictionary<int, int> subtreeHeight = new Dictionary<int, int>();
    private Dictionary<int, int> answer = new Dictionary<int, int>();
    
    public int[] TreeQueries(TreeNode root, int[] queries) {
        Dfs1(root, 0);
        Dfs2(root, 0);
        
        int[] result = new int[queries.Length];
        for (int i = 0; i < queries.Length; i++) {
            result[i] = answer[queries[i]];
        }
        return result;
    }
    
    private int Dfs1(TreeNode node, int depth) {
        if (node == null) return -1;
        
        nodeDepth[node.val] = depth;
        int leftHeight = Dfs1(node.left, depth + 1);
        int rightHeight = Dfs1(node.right, depth + 1);
        
        return subtreeHeight[node.val] = Math.Max(leftHeight, rightHeight) + 1;
    }
    
    private void Dfs2(TreeNode node, int maxDepthFromParent) {
        if (node == null) return;
        
        answer[node.val] = maxDepthFromParent;
        
        int leftHeight = node.left != null ? subtreeHeight[node.left.val] : -1;
        int rightHeight = node.right != null ? subtreeHeight[node.right.val] : -1;
        
        Dfs2(node.left, Math.Max(maxDepthFromParent, nodeDepth[node.val] + 1 + rightHeight));
        Dfs2(node.right, Math.Max(maxDepthFromParent, nodeDepth[node.val] + 1 + leftHeight));
    }
}

var treeQueries = function(root, queries) {
    const nodeDepth = new Map();
    const subtreeHeight = new Map();
    const answer = new Map();
    
    function dfs1(node, depth) {
        if (!node) return -1;
        
        nodeDepth.set(node.val, depth);
        const leftHeight = dfs1(node.left, depth + 1);
        const rightHeight = dfs1(node.right, depth + 1);
        
        const height = Math.max(leftHeight, rightHeight) + 1;
        subtreeHeight.set(node.val, height);
        return height;
    }
    
    function dfs2(node, maxDepthFromParent) {
        if (!node) return;
        
        answer.set(node.val, maxDepthFromParent);
        
        const leftHeight = node.left ? subtreeHeight.get(node.left.val) : -1;
        const rightHeight = node.right ? subtreeHeight.get(node.right.val) : -1;
        
        dfs2(node.left, Math.max(maxDepthFromParent, nodeDepth.get(node.val) + 1 + rightHeight));
        dfs2(node.right, Math.max(maxDepthFromParent, nodeDepth.get(node.val) + 1 + leftHeight));
    }
    
    dfs1(root, 0);
    dfs2(root, 0);
    
    return queries.map(query => answer.get(query));
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型复杂度
时间复杂度O(n + m)
空间复杂度O(n)

说明：n是树中节点数，m是查询数。预处理需要O(n)时间遍历整棵树两次，每个查询O(1)时间响应。空间复杂度主要用于存储哈希表和递归栈空间。

相关题目