# Height of Binary Tree After Subtree Removal Queries
**Difficulty:** HARD
[External](https://leetcode.com/problems/height-of-binary-tree-after-subtree-removal-queries)
Canonical: https://scaleengineer.com/dsa/problems/height-of-binary-tree-after-subtree-removal-queries
**Algorithms:** [Depth-First Search](https://scaleengineer.com/algorithms/depth-first-search), [Breadth-First Search](https://scaleengineer.com/algorithms/breadth-first-search)
**Data structures:** Array, Tree, Binary Tree
---
## Problem
You are given the `root` of a **binary tree** with `n` nodes. Each node is assigned a unique value from `1` to `n`. You are also given an array `queries` of size `m`.

You have to perform `m` **independent** queries on the tree where in the `ith` query you do the following:

* **Remove** the subtree rooted at the node with the value `queries[i]` from the tree. It is **guaranteed** that `queries[i]` will **not** be equal to the value of the root.

Return _an array_ `answer` _of size_ `m` _where_ `answer[i]` _is the height of the tree after performing the_ `ith` _query_.

**Note**:

* The queries are independent, so the tree returns to its **initial** state after each query.
* The height of a tree is the **number of edges in the longest simple path** from the root to some node in the tree.

**Example 1:**

![](https://assets.glich.co/dsa/height-of-binary-tree-after-subtree-removal-queries/image0.png) 

**Input:** root = [1,3,4,2,null,6,5,null,null,null,null,null,7], queries = [4]
**Output:** [2]
**Explanation:** The diagram above shows the tree after removing the subtree rooted at node with value 4.
The height of the tree is 2 (The path 1 -> 3 -> 2).

**Example 2:**

![](https://assets.glich.co/dsa/height-of-binary-tree-after-subtree-removal-queries/image1.png) 

**Input:** root = [5,8,9,2,1,3,7,4,6], queries = [3,2,4,8]
**Output:** [3,2,3,2]
**Explanation:** We have the following queries:
- Removing the subtree rooted at node with value 3. The height of the tree becomes 3 (The path 5 -> 8 -> 2 -> 4).
- Removing the subtree rooted at node with value 2. The height of the tree becomes 2 (The path 5 -> 8 -> 1).
- Removing the subtree rooted at node with value 4. The height of the tree becomes 3 (The path 5 -> 8 -> 2 -> 6).
- Removing the subtree rooted at node with value 8. The height of the tree becomes 2 (The path 5 -> 9 -> 3).

**Constraints:**

* The number of nodes in the tree is `n`.
* `2 <= n <= 105`
* `1 <= Node.val <= n`
* All the values in the tree are **unique**.
* `m == queries.length`
* `1 <= m <= min(n, 104)`
* `1 <= queries[i] <= n`
* `queries[i] != root.val`

# Approaches
## Brute Force Simulation
This approach directly simulates the process described in the problem. For each query, it finds the node to be removed, temporarily removes its subtree by modifying the parent's child pointer, calculates the height of the resulting tree, and then restores the tree to its original state for the next query.
**Time:** O(m * n) - For each of the `m` queries, we perform two traversals of the tree: one to find the node and its parent (O(n)), and another to calculate the height (O(n)). · **Space:** O(n) - The space is dominated by the recursion stack for the DFS traversals. In the worst case of a skewed tree, the recursion depth can be `n`.
**Pros:** Simple to understand and implement.; Directly follows the logic described in the problem statement.
**Cons:** Very inefficient due to repeated computations.; The time complexity of O(m * n) will cause a 'Time Limit Exceeded' (TLE) error for the given constraints.
### Explanation
The main idea is to iterate through each query one by one. For a given query value, we first need to locate the corresponding node `toRemove` and its `parent` in the tree. This can be done with a traversal like Depth First Search (DFS). Once the node and its parent are found, we determine if `toRemove` is a left or right child of `parent`. We then set the corresponding child pointer of `parent` to `null`, effectively removing the subtree.

After removal, we calculate the height of the modified tree. A standard recursive DFS function can compute the height of a tree by finding the maximum depth. The height is the number of edges on the longest path from the root to a leaf.

After recording the height, we must restore the tree to its initial state since the queries are independent. This is done by re-attaching `toRemove` to its `parent`. This entire process is repeated for all queries.

```java
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    // Helper to find the parent of a node with a given value.
    private TreeNode findParent(TreeNode root, int val) {
        if (root == null) return null;
        if ((root.left != null && root.left.val == val) || (root.right != null && root.right.val == val)) {
            return root;
        }
        TreeNode leftResult = findParent(root.left, val);
        if (leftResult != null) return leftResult;
        return findParent(root.right, val);
    }

    // Helper to calculate the height of a tree.
    private int calculateHeight(TreeNode root) {
        if (root == null) return -1; // Height of an empty tree is -1
        return 1 + Math.max(calculateHeight(root.left), calculateHeight(root.right));
    }

    public int[] treeQueries(TreeNode root, int[] queries) {
        int m = queries.length;
        int[] answer = new int[m];

        for (int i = 0; i < m; i++) {
            int queryVal = queries[i];
            // It's guaranteed that the query is not the root, so parent is not null.
            TreeNode parent = findParent(root, queryVal);
            TreeNode toRemove = null;
            boolean isLeftChild = false;

            if (parent.left != null && parent.left.val == queryVal) {
                toRemove = parent.left;
                isLeftChild = true;
                parent.left = null;
            } else {
                toRemove = parent.right;
                isLeftChild = false;
                parent.right = null;
            }

            answer[i] = calculateHeight(root);

            // Restore the tree for the next independent query
            if (isLeftChild) {
                parent.left = toRemove;
            } else {
                parent.right = toRemove;
            }
        }
        return answer;
    }
}
```
### Algorithm
- Initialize an answer array `ans` of the same size as `queries`.
- For each `query_val` in `queries`:
  - Find the node `toRemove` with value `query_val` and its `parent`. This requires a tree traversal (e.g., DFS).
  - Store which child of `parent` `toRemove` is (left or right).
  - Set the corresponding child pointer of `parent` to `null`.
  - Calculate the height of the current tree using a recursive height function. Store this in `ans`.
  - Restore the tree by setting the `parent`'s child pointer back to `toRemove`.
- Return `ans`.

## Optimal Two-Pass DFS Pre-computation
This approach avoids re-computation by pre-calculating the answer for every possible node removal in the tree. It uses two Depth First Search (DFS) traversals. The first pass gathers information about each node's depth and the height of its subtree. The second pass uses this information to efficiently calculate the height of the tree if that node's subtree were to be removed.
**Time:** O(n + m) - The first DFS pass takes O(n). The second DFS pass also takes O(n). Processing the `m` queries takes O(m) time for map lookups. · **Space:** O(n) - We use maps to store `depth`, `height`, and `answer` for `n` nodes, taking O(n) space. The recursion stack depth can also be up to O(n) in the worst case for a skewed tree.
**Pros:** Highly efficient with a linear time complexity.; Solves the problem within the given time constraints.; Answers all possible queries with one-time pre-computation.
**Cons:** More complex to understand and implement compared to the brute-force approach.; Requires careful handling of information passed between DFS calls.; Uses extra space for storing pre-computed values.
### Explanation
The core idea is that for any node `u`, the height of the tree after removing its subtree is the maximum depth among all leaves *not* in `u`'s subtree. We can pre-compute this value for every node `u` in the tree.

**First DFS Pass (post-order):** We traverse the tree to compute two properties for each node `u`:
- `depth[u]`: The depth of node `u` (number of edges from the root). This is passed down during the traversal.
- `height[u]`: The height of the subtree rooted at `u`. This is computed bottom-up. `height[u] = 1 + max(height(u.left), height(u.right))`.

**Second DFS Pass (pre-order):** We traverse the tree again to compute the answer for each node `u`. The answer for removing `u` is the length of the longest path from the root that does not enter `u`'s subtree. This path can either be a path that avoids `u`'s parent entirely, or a path that goes through `u`'s parent but then into `u`'s sibling's subtree. We can pass this information down the tree. For a node `curr`, we pass a value `max_h_from_ancestors`. When we recurse to a child, say `left_child`, the new value passed will be the maximum of `max_h_from_ancestors` and the length of the longest path through `curr`'s right child.

After the two passes, we have the answer pre-computed for every node. We can then answer all queries in O(1) time each by looking up the results.

```java
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    private Map<Integer, Integer> heightMap = new HashMap<>();
    private Map<Integer, Integer> depthMap = new HashMap<>();
    private Map<Integer, Integer> resultMap = new HashMap<>();

    public int[] treeQueries(TreeNode root, int[] queries) {
        // First pass: calculate height and depth for each node
        getHeightAndDepth(root, 0);
        
        // Second pass: calculate the answer for each node removal
        // The initial max_h_ancestors is -1 (height of an empty tree)
        calculateAnswers(root, -1);

        int[] ans = new int[queries.length];
        for (int i = 0; i < queries.length; i++) {
            ans[i] = resultMap.get(queries[i]);
        }
        return ans;
    }

    private int getHeightAndDepth(TreeNode node, int d) {
        if (node == null) {
            return -1;
        }
        depthMap.put(node.val, d);
        int leftHeight = getHeightAndDepth(node.left, d + 1);
        int rightHeight = getHeightAndDepth(node.right, d + 1);
        int h = 1 + Math.max(leftHeight, rightHeight);
        heightMap.put(node.val, h);
        return h;
    }

    private void calculateAnswers(TreeNode node, int maxAncestorHeight) {
        if (node == null) {
            return;
        }
        
        // The answer for removing the current node's subtree is the height
        // of the longest path that doesn't go through this node.
        resultMap.put(node.val, maxAncestorHeight);

        int leftHeight = (node.left == null) ? -1 : heightMap.get(node.left.val);
        int rightHeight = (node.right == null) ? -1 : heightMap.get(node.right.val);
        
        // For the left child, the competing path is the one going down the right sibling's branch
        // or the one coming from the ancestors.
        int maxForLeft = Math.max(maxAncestorHeight, (node.right == null) ? -1 : depthMap.get(node.val) + 1 + rightHeight);
        calculateAnswers(node.left, maxForLeft);

        // For the right child, the competing path is the one going down the left sibling's branch
        // or the one coming from the ancestors.
        int maxForRight = Math.max(maxAncestorHeight, (node.left == null) ? -1 : depthMap.get(node.val) + 1 + leftHeight);
        calculateAnswers(node.right, maxForRight);
    }
}
```
### Algorithm
- Create maps to store `depth`, `height`, and the final `answer` for each node's value.
- **First DFS (post-order traversal):**
  - Define a function `dfs_height(node, d)`.
  - At each `node`, store its depth `d` in the `depth` map.
  - Recursively call for left and right children with depth `d+1`.
  - Compute `node`'s height based on its children's heights and store it in the `height` map.
  - Call `dfs_height(root, 0)`.
- **Second DFS (pre-order traversal):**
  - Define a function `dfs_answer(node, max_h_ancestors)`.
  - At each `node`, store `max_h_ancestors` as the answer for removing this node's subtree.
  - For the left child, the new `max_h_ancestors` will be `max(current_max_h_ancestors, longest_path_through_right_sibling)`.
  - For the right child, the new `max_h_ancestors` will be `max(current_max_h_ancestors, longest_path_through_left_sibling)`.
  - Recursively call `dfs_answer` for left and right children with their newly computed `max_h_ancestors`.
  - Call `dfs_answer(root, -1)`.
- **Process Queries:**
  - Initialize an answer array `res`.
  - For each `query_val` in `queries`, look up the pre-computed answer from the `answer` map and add it to `res`.
  - Return `res`.

# Solutions
### Java

```java
/** * 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 { private Map < TreeNode , Integer > d = new HashMap <>(); private int [] res ; public int [] treeQueries ( TreeNode root , int [] queries ) { f ( root ); res = new int [ d . size () + 1 ]; d . put ( null , 0 ); dfs ( root , - 1 , 0 ); int m = queries . length ; int [] ans = new int [ m ]; for ( int i = 0 ; i < m ; ++ i ) { ans [ i ] = res [ queries [ i ]]; } return ans ; } private void dfs ( TreeNode root , int depth , int rest ) { if ( root == null ) { return ; } ++ depth ; res [ root . val ] = rest ; dfs ( root . left , depth , Math . max ( rest , depth + d . get ( root . right ))); dfs ( root . right , depth , Math . max ( rest , depth + d . get ( root . left ))); } private int f ( TreeNode root ) { if ( root == null ) { return 0 ; } int l = f ( root . left ), r = f ( root . right ); d . put ( root , 1 + Math . max ( l , r )); return d . get ( root ); } }
```

### CPP

```cpp
/** * 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 { public: vector < int > treeQueries ( TreeNode * root , vector < int >& queries ) { unordered_map < TreeNode * , int > d ; function < int ( TreeNode * ) > f = [ & ]( TreeNode * root ) -> int { if ( ! root ) return 0 ; int l = f ( root -> left ), r = f ( root -> right ); d [ root ] = 1 + max ( l , r ); return d [ root ]; }; f ( root ); vector < int > res ( d . size () + 1 ); function < void ( TreeNode * , int , int ) > dfs = [ & ]( TreeNode * root , int depth , int rest ) { if ( ! root ) return ; ++ depth ; res [ root -> val ] = rest ; dfs ( root -> left , depth , max ( rest , depth + d [ root -> right ])); dfs ( root -> right , depth , max ( rest , depth + d [ root -> left ])); }; dfs ( root , - 1 , 0 ); vector < int > ans ; for ( int v : queries ) ans . emplace_back ( res [ v ]); return ans ; } };
```

### Python

```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 ]: def f ( root ): if root is None : return 0 l , r = f ( root . left ), f ( root . right ) d [ root ] = 1 + max ( l , r ) return d [ root ] def dfs ( root , depth , rest ): if root is None : return depth += 1 res [ root . val ] = rest dfs ( root . left , depth , max ( rest , depth + d [ root . right ])) dfs ( root . right , depth , max ( rest , depth + d [ root . left ])) d = defaultdict ( int ) f ( root ) res = [ 0 ] * ( len ( d ) + 1 ) dfs ( root , - 1 , 0 ) return [ res [ v ] for v in queries ]
```