python
# Pseudo-code for Brute Force
def solve_naive(root, queries):
    answers = []
    for q_val in queries:
        # 1. Mark node q_val as "removed"
        # 2. Run DFS from root to find max depth, ignoring "removed" paths
        current_height = calculate_height(root, ignore=q_val)
        answers.append(current_height)
    return answers

Why this fails: The time complexity is $O(M \times N)$, where $M$ is the number of queries and $N$ is the number of nodes. With constraints up to $10^5$ for both $N$ and $M$, this results in roughly $10^{10}$ operations, which inevitably causes a Time Limit Exceeded (TLE) error. We need a solution that answers each query in $O(1)$ time after preprocessing.

Core Insight for Height of Binary Tree After Subtree Removal Queries

The key intuition is to realize that the height of the tree is determined by the longest path from the root to a leaf. If we remove a subtree rooted at node u, we break all paths passing through u. The new height of the tree will be determined by the longest path that does not go through u.

Every path from the root to a leaf passes through exactly one node at any specific depth level $L$ (up to the path's end). Therefore, if we remove a node u at depth d, the new maximum height is determined by the "cousins" of u—other nodes residing at the same depth d.

Specifically, for any depth level, we only care about the nodes that produce the longest paths.

1. Preprocessing: We can compute the depth of every node and the height of the subtree rooted at every node.
2. The Invariant: If we remove node u at depth depth[u], the new tree height is the maximum of:
   • The longest path passing through a different node at the same depth depth[u].
   • The depth of the parent (if no other nodes exist at this depth).

By storing the two largest subtree heights for every depth level, we can answer queries instantly. If the node being removed holds the largest height at its level, the answer is the second largest. Otherwise, the answer remains the largest.

Visual Description: Imagine the recursion tree. As the Postorder Traversal bubbles up from the leaves, each node calculates its height as 1 + max(left_height, right_height). We essentially "annotate" every horizontal level of the tree with the subtree heights found at that level. When a query cuts a branch, we simply look horizontally at that level to see which other branch is longest.

Algorithm Strategy: Tree DFS - Recursive Postorder Traversal

1. Data Structures:

   • depth_map: Stores the depth (distance from root) of each node.
   • height_map: Stores the height (distance to deepest leaf) of the subtree rooted at each node.
   • level_maxes: A hash map (or array) where keys are depths. For each depth, we store the two largest subtree heights found at that depth.

2. Preprocessing (DFS):

   • Perform a DFS (Postorder logic).
   • For a node u at depth d:
     • Recursively calculate heights of left and right children.
     • height[u] = 1 + max(height[left], height[right]).
     • Record depth[u] = d.
     • Update level_maxes[d] with height[u]. Keep only the top 2 distinct values (associated with their node IDs to handle ties correctly).

3. Query Processing:

   • For each query node q:
     • Retrieve d = depth_map[q] and h = height_map[q].
     • Look at level_maxes[d].
     • If q is the node responsible for the maximum height at this level, the new tree height is d + second_max_height_at_level.
     • Otherwise, the new tree height is d + max_height_at_level.

Execution Flow

1. Initialization: Create maps for depths, heights, and level-wise maximums.
2. DFS Traversal (Root, Depth 0):
   • If node is null, return -1 (height of null).
   • Recursively call DFS for left child with depth + 1.
   • Recursively call DFS for right child with depth + 1.
   • Current subtree height = 1 + max(left_sub_height, right_sub_height).
   • Store depth and height for the current node.
   • Update the "Top 2" list for the current depth. If the current height is larger than the stored maximums, shift them down.
3. Answering Queries:
   • Loop through queries.
   • Identify the depth L of the query node.
   • Check the precomputed top heights for level L.
   • If the query node contributes the absolute longest path at level L, use the secondary longest path.
   • Calculate Result = L + selected_subtree_height.
   • Append to results.

Proof of Correctness

The height of the binary tree is defined as the maximum number of edges from the root to any leaf. Any path from the root to a leaf can be decomposed into two parts relative to a specific depth $L$: the path from Root to a node $v$ at depth $L$ (length $L$), and the path from $v$ to a leaf (length $height[v]$).

The global height is $\max_{\forall v} (depth[v] + height[v])$. Since we are removing a subtree at node $q$ (where $depth[q] = L$), we are effectively removing the term $(L + height[q])$ from the set of all path lengths. Because all paths reaching depth greater than $L$ must pass through some node at depth $L$, the new maximum height must be formed by a path passing through a sibling or cousin of $q$ at depth $L$. Therefore, $\max_{v \in Level_L, v \neq q} (L + height[v])$ gives the correct new height.

Reference Implementation

C++ Solution for LeetCode 2458

cpp
#include <vector>
#include <unordered_map>
#include <algorithm>
#include <iostream>

using namespace std;

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
    // Map to store the depth of each node: node_val -> depth
    unordered_map<int, int> depthMap;
    // Map to store the height of the subtree rooted at each node: node_val -> height
    unordered_map<int, int> heightMap;
    // Stores the top 2 largest heights for each depth level
    // levelMaxes[d] = { {height1, node1}, {height2, node2} }
    unordered_map<int, vector<pair<int, int>>> levelMaxes;

    int computeHeights(TreeNode* node, int depth) {
        if (!node) return -1;

        depthMap[node->val] = depth;

        int leftH = computeHeights(node->left, depth + 1);
        int rightH = computeHeights(node->right, depth + 1);

        int currentH = 1 + max(leftH, rightH);
        heightMap[node->val] = currentH;

        // Update top 2 heights for this depth
        auto& top2 = levelMaxes[depth];
        top2.push_back({currentH, node->val});
        sort(top2.rbegin(), top2.rend()); // Sort descending
        if (top2.size() > 2) {
            top2.pop_back(); // Keep only top 2
        }

        return currentH;
    }

public:
    vector<int> treeQueries(TreeNode* root, vector<int>& queries) {
        depthMap.clear();
        heightMap.clear();
        levelMaxes.clear();

        // 1. Precompute depths and subtree heights using Postorder DFS
        computeHeights(root, 0);

        vector<int> answer;
        answer.reserve(queries.size());

        // 2. Answer queries in O(1)
        for (int q : queries) {
            int d = depthMap[q];
            int h = heightMap[q];
            
            // Retrieve top heights at this depth
            auto& tops = levelMaxes[d];
            
            // If the node being removed is the one with the largest height at this level
            if (tops[0].second == q) {
                if (tops.size() == 2) {
                    // Use the second largest height
                    answer.push_back(d + tops[1].first);
                } else {
                    // No other node at this depth extends further; max path ends at level d-1
                    answer.push_back(d - 1);
                }
            } else {
                // The node removed is not the largest, so global max remains unchanged
                answer.push_back(d + tops[0].first);
            }
        }

        return answer;
    }
};

Java Solution for LeetCode 2458

java
import java.util.*;

/**
 * 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 {
    // node val -> depth
    private Map<Integer, Integer> depthMap = new HashMap<>();
    // node val -> subtree height
    private Map<Integer, Integer> heightMap = new HashMap<>();
    // depth -> list of top 2 pairs (height, nodeVal)
    private Map<Integer, List<int[]>> levelMaxes = new HashMap<>();

    public int[] treeQueries(TreeNode root, int[] queries) {
        // 1. Precompute depths and heights
        computeHeights(root, 0);

        int[] answer = new int[queries.length];
        
        // 2. Process queries
        for (int i = 0; i < queries.length; i++) {
            int q = queries[i];
            int d = depthMap.get(q);
            
            List<int[]> tops = levelMaxes.get(d);
            
            // Check if the removed node holds the max height at this level
            if (tops.get(0)[1] == q) {
                if (tops.size() > 1) {
                    answer[i] = d + tops.get(1)[0];
                } else {
                    // No cousins, path stops at parent
                    answer[i] = d - 1;
                }
            } else {
                // Max height is determined by a different node
                answer[i] = d + tops.get(0)[0];
            }
        }
        
        return answer;
    }

    private int computeHeights(TreeNode node, int depth) {
        if (node == null) return -1;

        depthMap.put(node.val, depth);

        int leftH = computeHeights(node.left, depth + 1);
        int rightH = computeHeights(node.right, depth + 1);

        int currentH = 1 + Math.max(leftH, rightH);
        heightMap.put(node.val, currentH);

        // Manage top 2 heights for this level
        levelMaxes.putIfAbsent(depth, new ArrayList<>());
        List<int[]> tops = levelMaxes.get(depth);
        
        tops.add(new int[]{currentH, node.val});
        // Sort descending by height
        tops.sort((a, b) -> b[0] - a[0]);
        
        if (tops.size() > 2) {
            tops.remove(2);
        }

        return currentH;
    }
}

Python Solution for LeetCode 2458

python
# 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 treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
        depth_map = {}
        height_map = {}
        # Stores top 2 heights at each level: depth -> [(height, node_val), ...]
        level_maxes = collections.defaultdict(list)

        def compute_heights(node, depth):
            if not node:
                return -1
            
            depth_map[node.val] = depth
            
            left_h = compute_heights(node.left, depth + 1)
            right_h = compute_heights(node.right, depth + 1)
            
            current_h = 1 + max(left_h, right_h)
            height_map[node.val] = current_h
            
            # Update top 2 for this level
            level_maxes[depth].append((current_h, node.val))
            level_maxes[depth].sort(key=lambda x: x[0], reverse=True)
            if len(level_maxes[depth]) > 2:
                level_maxes[depth].pop()
                
            return current_h

        # 1. Precompute
        compute_heights(root, 0)
        
        answers = []
        
        # 2. Answer queries
        for q in queries:
            d = depth_map[q]
            tops = level_maxes[d]
            
            # If q is the node with the largest height at this level
            if tops[0][1] == q:
                if len(tops) > 1:
                    answers.append(d + tops[1][0])
                else:
                    answers.append(d - 1)
            else:
                answers.append(d + tops[0][0])
                
        return answers