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

思路

给定一棵二叉树的根节点 root 和一个整数数组 queries，对于每个查询 queries[i]，返回移除以 queries[i] 为根节点的子树后，剩余树的高度。

解题过程

为了求解这个问题，我们采用深度优先搜索（DFS）策略，同时记录每个节点的访问顺序、深度、子树大小，并计算每个节点左侧和右侧的最大深度。这样，当我们移除某个子树时，我们可以快速计算出剩余树的最大深度。

复杂度

• 时间复杂度: $O(n)$，其中 nn 是二叉树中节点的数量。DFS 遍历需要访问每个节点一次，计算最大深度和处理查询也只需要线性时间。
• 空间复杂度: $O(n)$，主要由辅助数组占用的空间决定。

Code

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public static final int MAXN = 100010;
    public static int[] dfn = new int[MAXN];
    public static int[] size = new int[MAXN];
    public static int[] deep = new int[MAXN];
    public static int[] maxl = new int[MAXN];
    public static int[] maxr = new int[MAXN];
    public static int dfnCnt;

    public static int[] treeQueries(TreeNode root, int[] queries) {
        dfnCnt = 0;
        f(root, 0);
        for(int i = 1; i <= dfnCnt; i++) {
            maxl[i] = Math.max(maxl[i - 1], deep[i]);
        }
        maxr[dfnCnt + 1] = 0;
        for(int i = dfnCnt; i >= 1; i--) {
            maxr[i] = Math.max(maxr[i + 1], deep[i]);
        }
        int m = queries.length;
        int[] ans = new int[m];
        for(int i = 0; i < m; i++) {
            int leftMax = maxl[dfn[queries[i]] - 1];
            int rightMax = maxr[dfn[queries[i]] + size[dfn[queries[i]]]];
            ans[i] = Math.max(leftMax, rightMax);
        }
        return ans;

    }
    public static void f(TreeNode x, int k) {
        int i = ++dfnCnt;
        dfn[x.val] = i;
        deep[i] = k;
        size[i] = 1;
        if(x.left != null) {
            f(x.left, k + 1);
            size[i] += size[dfn[x.left.val]];
        }
        if(x.right != null) {
            f(x.right, k + 1);
            size[i] += size[dfn[x.right.val]];
        }
    }
}