2458. Height of Binary Tree After Subtree Removal Queries - Detailed Explanation

Problem Statement

You are given a binary tree in which each node has a unique value. You are also given an array of queries. Each query is a value corresponding to a node in the tree. For each query, you need to determine the height of the tree after removing the subtree rooted at the given node (i.e., deleting that node and all of its descendants).

Definitions:

• Binary Tree Height:
  The height of a binary tree is defined as the number of edges on the longest downward path from the root to a leaf. For example, a tree consisting of a single node (the root) has a height of 0.

• Subtree Removal:
  Removing the subtree rooted at a node means that the node and all nodes in its subtree are deleted from the original tree. After this removal, the remaining tree is re-evaluated to determine its new height.

Input:

1. Tree Root:
   The root of a binary tree (where each node contains a unique integer value).

2. Queries:
   An array of integers, where each integer represents the value of a node whose subtree should be removed. The queries are independent; each query considers the removal from the original tree (not cumulatively).

Output:

For each query, return the height of the tree after removing the subtree rooted at the queried node. The answers should be returned in an array where each element corresponds to a query.

Example

Consider the following binary tree:

        1
       / \
      2   3
     /
    4

• Original Tree Height:
  The longest path is from 1 → 2 → 4, which contains 2 edges, so the height is 2.

• Queries: Suppose we have queries [3, 2, 4].

1. Query = 3 (Remove subtree rooted at node 3):

   • Remaining Tree:

         1
        /
       2
      /
     4
     

   • New Height: The longest path is 1 → 2 → 4 (2 edges), so the height remains 2.

2. Query = 2 (Remove subtree rooted at node 2):

   • Remaining Tree:

         1
          \
           3
     

   • New Height: The longest path is 1 → 3 (1 edge), so the height becomes 1.

3. Query = 4 (Remove subtree rooted at node 4):

   • Remaining Tree:

         1
        / \
       2   3
     

   • New Height: The longest path is either 1 → 2 or 1 → 3 (1 edge), so the height becomes 1.

• Output:
  For queries [3, 2, 4], the output should be [2, 1, 1].

Constraints

• The tree contains at least one node.
• Each node's value is unique.
• The number of queries can be large, so an efficient precomputation strategy is required to answer each query in O(1) time after an initial O(n) preprocessing.

This problem challenges you to combine information from both the child subtrees and the path from the root (the “up‑path”) to answer subtree removal queries efficiently. A standard solution uses two DFS traversals:

1. DFS₁ (Bottom‐Up):
   Traverse the tree from the root and compute for every node:

   • depth[node] – the number of edges from the root to that node (with root’s depth = 0).
   • height[node] – the height of the subtree rooted at that node (i.e. the number of edges on the longest downward path from the node; a leaf’s height is 0).

2. DFS₂ (Top‑Down “Up‑Values”):
   Now, for every node, compute an “up‑value” (let’s call it up[node]) that represents the maximum number of edges on a path from the root that does not go into that node’s subtree. In other words, if you were to remove the subtree rooted at that node, then up[node] is the height of the remaining tree. When going from a parent to a child the new up‑value is determined by:

   • The up‑value coming from the parent (i.e. the best path “from above”) plus one (for the edge from parent to child), and
   • The height of the sibling subtree (if one exists) plus the parent’s depth plus one (because the best alternative path might go from the parent into its other child).

   In pseudocode for a child u of parent p:

   if u is the left child:
       candidate = (if p.right exists then depth[p] + 1 + height[p.right] else 0)
   if u is the right child:
       candidate = (if p.left exists then depth[p] + 1 + height[p.left] else 0)
   newUp = max( up[p] + 1, candidate )
   up[u] = newUp
   

   (The exact details are subtle; note that the definitions assume “depth” and “height” are measured in edges. Also, when a candidate does not exist we treat it as 0.)

After both DFS passes are complete, for every node (identified by its unique value) we have precomputed ans[node] = up[node]—which is exactly the height of the tree after removing the subtree rooted at that node. Then, each query (given as a node’s value) is answered in O(1) by looking up ans[node].

Step‑by‑Step Walkthrough (High‑Level)

1. Compute depth and subtree height (DFS₁):
   – Start at the root with depth 0.
   – For each node, recursively compute the height of its left and right subtrees; then set
   height[node] = max(height(left), height(right)) + 1
   (with a leaf’s height defined as 0). Also record depth[node] along the way.

2. Compute “up‑values” (DFS₂):
   – For the root, define up[root] = 0 (if you remove the entire tree, the “height” is 0).
   – When moving from a parent p (with known up[p]) to a child u, determine the best alternative path that does not go into u’s subtree. If u is the left child, then an alternative is the right child of p (if it exists); if u is the right child, an alternative is the left child. In each case add the extra edge from p and use p’s depth to “lift” the candidate value. Then, set
   up[u] = max( up[p] + 1, candidate )
   – Recurse on both children.

3. Answer Queries:
   – For each query (a node value q), return up[q] as the height of the tree after removal of q’s subtree.

Complexity Analysis

• Time Complexity:
  – DFS₁ visits each node once and does O(1) work per node.
  – DFS₂ also visits each node once.
  – Answering each query is O(1).
  Overall, if there are n nodes and m queries, the total time is O(n + m).

• Space Complexity:
  – We store a “depth” and “height” value for each node, plus the up‑values in a hash map (or array) – O(n) extra space.
  – The recursion stack uses O(n) in the worst case.

Common Pitfalls and Edge Cases

• Pitfall – Incorrect Definitions:
  Make sure you are consistent in how you define “depth” and “height” (whether in terms of edges or nodes). In this explanation we use “depth” as the number of edges from the root (so the root has depth 0) and “height” as the number of edges on the longest downward path (a leaf’s height is 0). Adjust the recurrence accordingly if you use 1-indexed definitions.

• Pitfall – Forgetting to Backtrack:
  In DFS₂ (the top‑down pass) you must correctly combine the “up‑value” from the parent with the candidate from the sibling’s subtree.

• Edge Case – Removal at the Root:
  Removing the subtree at the root removes the entire tree so the answer is defined to be 0.

• Edge Case – When a Sibling Does Not Exist:
  If the node has no sibling (i.e. if it is the only child), then the candidate from that side is taken as 0.

Pseudocode Summary

function dfs1(node, d):
    if node is null: return -1
    depth[node] = d
    leftH = dfs1(node.left, d+1)
    rightH = dfs1(node.right, d+1)
    height[node] = max(leftH, rightH) + 1
    return height[node]

function dfs2(node, up):
    ans[node.val] = up
    if node.left exists:
        candidate = (node.right exists ? depth[node] + 1 + height[node.right] : 0)
        newUp = max(up + 1, candidate)
        dfs2(node.left, newUp)
    if node.right exists:
        candidate = (node.left exists ? depth[node] + 1 + height[node.left] : 0)
        newUp = max(up + 1, candidate)
        dfs2(node.right, newUp)

Then for each query on a node with value q, answer = ans[q].

Code Implementations

Below are sample implementations in both Python and Java. (Note that in an actual solution the tree is given via a TreeNode class with unique values; the queries are given as an array of integers. The solution precomputes answers for every node.)

Python Implementation