Leetcode 2458. 移除子树后的二叉树高度

给你一棵 二叉树 的根节点 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
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
36
37
38
39
40
41
42
43
44
45

• 首先统计出所有子树的高度。然后再计算出移除每一个结点后树的高度变化，对于某一层来说，具有高度最大和次大的结点，当移除非高度最大的节点时，树的最大高度不会改变，为这一层节点最大高度+根节点到该层的路径边数；当移除最大高度时，树的最大高度变为次大高度+根节点到该层的路径边数。当一层只有一个节点时，根节点到该层的路径数要-1。
• 时间复杂度：$O(n)$
• 空间复杂度：$O(n)$

class Solution {
    Map<Integer, Integer> height = new HashMap<>();
    Map<Integer, Integer> valToAns = new HashMap<>();
    public int[] treeQueries(TreeNode root, int[] queries) {
        int n = queries.length;
        int[] ans = new int[n];
        getLength(root); 
        Queue<TreeNode> q = new LinkedList<>();
        List<Integer> vals = new ArrayList<>();
        q.add(root);  
        int step = 0;
        while (!q.isEmpty()) {
            int size = q.size();
            int maxv = 0, maxv2 = -1;
            for (int i = 1; i <= size; i++) {
                TreeNode p = q.poll();
                vals.add(p.val);
                if (height.get(p.val) > maxv) {
                    maxv2 = maxv;
                    maxv = height.get(p.val);
                } else if (height.get(p.val) > maxv2) {
                    maxv2 = height.get(p.val);
                }
                if (p.left != null) q.add(p.left); 
                if (p.right != null) q.add(p.right);
            }
            for (int i = 0; i < size; i++){
                int val = vals.get(i);
                if (height.get(val) == maxv) valToAns.put(val, maxv2 + step);
                else valToAns.put(val, maxv + step);
                if (size == 1) valToAns.put(val, valToAns.get(val) - 1); //特殊情况
            }
            vals.clear();
            step++;
        } 
        for (int i = 0; i < n; i++) ans[i] = valToAns.get(queries[i]);
        return ans;
    }
    int getLength(TreeNode root) {
        if (root == null) return -1;    
        int l = getLength(root.left);            
        int r = getLength(root.right);        
        int len = Math.max(l, r) + 1;  
        height.put(root.val, len);         
        return len;         
    }
}
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
36
37
38
39
40
41
42
43
44
45
46
47

• 使用dfs序来给每一个节点进行标号,dfs即dfs第一次搜索到该节点的次序，通过该种方式，可以发现以根节点出发，访问所有子节点的次序必然是连续的，那么可以统计该节点为根的子树能够包括的范围。使用depth[] 来统计每个以dfs序为索引的节点的深度，如5:d[1] = 0。
• 当删除某个节点，如2，等于删除以2为根的子树所包含的所有范围即$[3,5]$, 那么我们只需要在$[1,2]$和$[6,9]$统计每个节点的最大深度。可以使用前缀最大值和后缀最大值快速求出$[1,l-1]$和$[r+1,n]$的最大深度进行比较。

• 时间复杂度：$O(n)$
• 空间复杂度：$O(n)$

class Solution {
    int N = 100005;
    int[] l = new int[N], r = new int[N], depth = new int[N], lmax = new int[N], rmax= new int[N];
    int cnt = 0;
    public int[] treeQueries(TreeNode root, int[] queries) {
        dfs(root, 0);
        int[] ans = new int[queries.length];
        for (int i = 1; i <= cnt; i++) lmax[i] = Math.max(depth[i], lmax[i - 1]);
        for (int i = cnt; i >= 0; i--) rmax[i] = Math.max(depth[i], rmax[i + 1]);
        for (int i = 0; i < queries.length; i++) ans[i] = Math.max(lmax[l[queries[i]] - 1],  rmax[r[queries[i]] + 1]);
        return ans;
    }
    void dfs(TreeNode root, int level) {
        if (root == null) return;
        depth[++cnt] = level;
        l[root.val] = cnt;
        dfs(root.left, level + 1);
        dfs(root.right, level + 1);
        r[root.val] = cnt;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21