2458. Height of Binary Tree After Subtree Removal Queries

HardTreeDepth-First SearchBreadth-First SearchArrayBinary Tree

Problem Description

You have a binary tree with n nodes, where each node has a unique value from 1 to n. You're given an array queries of size m.

For each query queries[i], you need to:

Remove the subtree rooted at the node with value queries[i] from the tree
Calculate the height of the remaining tree
Note that queries[i] will never be the root node

The key point is that each query is independent - the tree returns to its original state before processing the next query.

The height of a tree is defined as the number of edges in the longest path from the root to any leaf node.

Your task is to return an array answer where answer[i] represents the height of the tree after removing the subtree specified by queries[i].

For example, if you have a tree and remove a subtree rooted at node x, you need to find the maximum depth among all remaining paths from the root. This becomes the new height of the tree for that particular query.

The solution uses two DFS traversals:

First DFS (f function): Calculates the depth of each subtree and stores it in dictionary d, where d[node] represents the height of the subtree rooted at that node
Second DFS (dfs function): For each node, calculates what the tree height would be if that node's subtree were removed, considering the maximum between:
- The current path depth up to that node's parent
- The depth through the sibling subtree (if it exists)

The res array stores the final height for each node value when that node's subtree is removed, which can then be used to answer all queries efficiently.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

Yes: A binary tree is a special type of graph where each node has at most two children and there are no cycles.

Is it a tree?

Yes: The problem explicitly states we're working with a binary tree structure with a root node and parent-child relationships.

DFS

Yes: This leads us to DFS (Depth-First Search) as the recommended approach.

Conclusion: The flowchart suggests using DFS for this tree problem.

The DFS pattern is particularly suitable here because:

We need to traverse the entire tree to calculate subtree heights
We need to explore each path from root to leaves to determine the maximum depth
The problem requires processing each node and its subtrees recursively
We need to collect information (heights) from child nodes before processing parent nodes

The solution implements two DFS traversals:

First DFS: Post-order traversal to calculate the height of each subtree
Second DFS: Pre-order traversal to determine what the tree height would be if each node's subtree were removed

This two-pass DFS approach efficiently solves the problem by precomputing the necessary information to answer all queries in O(1) time each.

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

Intuition

When we remove a subtree from the tree, we're essentially cutting off an entire branch. The key insight is that the new tree height will be determined by the longest remaining path from the root.

Think about what happens when we remove a subtree rooted at node x:

All nodes in x's subtree disappear
The rest of the tree remains intact
The new height is the maximum depth among all remaining paths

The naive approach would be to actually remove each queried subtree and recalculate the height, but this would be inefficient for multiple queries. Instead, we can precompute what the height would be if any node were removed.

Here's the clever observation: When we remove a node's subtree, the new tree height is the maximum of:

The depth of paths that don't go through that node at all
The depth from root to that node's parent (since we stop there after removal)

To efficiently calculate this, we need two pieces of information:

The height of each subtree (calculated in the first DFS)
For each node, what would be the maximum depth if we removed it

During the second DFS traversal, as we visit each node, we maintain rest - the maximum height achievable without using the current node's subtree. This includes:

Paths we've already explored (from ancestors and their other children)
The sibling subtree's contribution (if current node has a sibling)

For example, when we're at a left child, rest would consider:

The current depth to this node's parent
The maximum depth achievable through the right sibling subtree

This way, we precompute for every possible removal in just two tree traversals, making each query answerable in O(1) time by simply looking up the precomputed value.

Learn more about Tree, Depth-First Search, Breadth-First Search and Binary Tree patterns.

Solution Implementation

1# Definition for a binary tree node. 2# class TreeNode: 3# def __init__(self, val=0, left=None, right=None): 4# self.val = val 5# self.left = left 6# self.right = right 7 8from typing import Optional, List 9from collections import defaultdict 10 11class Solution: 12 def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]: 13 def calculate_height(node: Optional[TreeNode]) -> int: 14 """ 15 Calculate and store the height of each subtree. 16 Height is defined as the number of nodes on the longest path from node to leaf. 17 """ 18 if node is None: 19 return 0 20 21 left_height = calculate_height(node.left) 22 right_height = calculate_height(node.right) 23 24 # Store height of current subtree (1 + max height of children) 25 node_heights[node] = 1 + max(left_height, right_height) 26 27 return node_heights[node] 28 29 def compute_max_depth_without_subtree(node: Optional[TreeNode], current_depth: int, max_depth_without_current: int) -> None: 30 """ 31 For each node, compute the maximum depth of the tree if that node's subtree is removed. 32 33 Args: 34 node: Current node being processed 35 current_depth: Depth of the current node in the tree 36 max_depth_without_current: Maximum depth achievable without the current subtree 37 """ 38 if node is None: 39 return 40 41 # Move to the next depth level 42 current_depth += 1 43 44 # Store the maximum depth if this node's subtree is removed 45 max_depth_after_removal[node.val] = max_depth_without_current 46 47 # Process left child: if left subtree is removed, we can still reach depths through right subtree 48 compute_max_depth_without_subtree( 49 node.left, 50 current_depth, 51 max(max_depth_without_current, current_depth + node_heights[node.right]) 52 ) 53 54 # Process right child: if right subtree is removed, we can still reach depths through left subtree 55 compute_max_depth_without_subtree( 56 node.right, 57 current_depth, 58 max(max_depth_without_current, current_depth + node_heights[node.left]) 59 ) 60 61 # Dictionary to store heights of each subtree 62 node_heights = defaultdict(int) 63 64 # First pass: calculate heights of all subtrees 65 calculate_height(root) 66 67 # Array to store the result for each node value 68 # Size is number of nodes + 1 to handle 1-indexed node values 69 max_depth_after_removal = [0] * (len(node_heights) + 1) 70 71 # Second pass: compute maximum depth after removing each node's subtree 72 compute_max_depth_without_subtree(root, -1, 0) 73 74 # Return results for each query 75 return [max_depth_after_removal[node_val] for node_val in queries] 76

1/** 2 * Definition for a binary tree node. 3 * public class TreeNode { 4 * int val; 5 * TreeNode left; 6 * TreeNode right; 7 * TreeNode() {} 8 * TreeNode(int val) { this.val = val; } 9 * TreeNode(int val, TreeNode left, TreeNode right) { 10 * this.val = val; 11 * this.left = left; 12 * this.right = right; 13 * } 14 * } 15 */ 16class Solution { 17 // Map to store the height of each node's subtree 18 private Map<TreeNode, Integer> nodeHeightMap = new HashMap<>(); 19 20 // Array to store the maximum height after removing each node 21 private int[] maxHeightAfterRemoval; 22 23 /** 24 * Main method to process tree queries 25 * @param root The root of the binary tree 26 * @param queries Array of node values to query 27 * @return Array of maximum tree heights after removing each queried node 28 */ 29 public int[] treeQueries(TreeNode root, int[] queries) { 30 // Calculate and store the height of each subtree 31 calculateSubtreeHeights(root); 32 33 // Initialize result array (size is tree size + 1 for 1-indexed node values) 34 maxHeightAfterRemoval = new int[nodeHeightMap.size() + 1]; 35 36 // Add null node with height 0 for handling leaf nodes 37 nodeHeightMap.put(null, 0); 38 39 // Calculate maximum possible height when each node is removed 40 calculateMaxHeightsAfterRemoval(root, -1, 0); 41 42 // Build answer array based on queries 43 int queryCount = queries.length; 44 int[] answer = new int[queryCount]; 45 for (int i = 0; i < queryCount; i++) { 46 answer[i] = maxHeightAfterRemoval[queries[i]]; 47 } 48 49 return answer; 50 } 51 52 /** 53 * DFS to calculate the maximum height of the tree when each node is removed 54 * @param currentNode Current node being processed 55 * @param currentDepth Current depth in the tree 56 * @param maxHeightFromOtherPaths Maximum height achievable from alternative paths 57 */ 58 private void calculateMaxHeightsAfterRemoval(TreeNode currentNode, int currentDepth, 59 int maxHeightFromOtherPaths) { 60 if (currentNode == null) { 61 return; 62 } 63 64 // Increment depth for current node 65 currentDepth++; 66 67 // Store the maximum height when this node is removed 68 maxHeightAfterRemoval[currentNode.val] = maxHeightFromOtherPaths; 69 70 // Process left child: if left child is removed, consider right subtree height 71 int maxHeightIfLeftRemoved = Math.max(maxHeightFromOtherPaths, 72 currentDepth + nodeHeightMap.get(currentNode.right)); 73 calculateMaxHeightsAfterRemoval(currentNode.left, currentDepth, maxHeightIfLeftRemoved); 74 75 // Process right child: if right child is removed, consider left subtree height 76 int maxHeightIfRightRemoved = Math.max(maxHeightFromOtherPaths, 77 currentDepth + nodeHeightMap.get(currentNode.left)); 78 calculateMaxHeightsAfterRemoval(currentNode.right, currentDepth, maxHeightIfRightRemoved); 79 } 80 81 /** 82 * Calculate the height of each subtree rooted at each node 83 * @param currentNode Current node being processed 84 * @return Height of the subtree rooted at currentNode 85 */ 86 private int calculateSubtreeHeights(TreeNode currentNode) { 87 if (currentNode == null) { 88 return 0; 89 } 90 91 // Recursively calculate heights of left and right subtrees 92 int leftSubtreeHeight = calculateSubtreeHeights(currentNode.left); 93 int rightSubtreeHeight = calculateSubtreeHeights(currentNode.right); 94 95 // Height of current subtree is 1 + max of children heights 96 int currentSubtreeHeight = 1 + Math.max(leftSubtreeHeight, rightSubtreeHeight); 97 nodeHeightMap.put(currentNode, currentSubtreeHeight); 98 99 return currentSubtreeHeight; 100 } 101} 102

1/** 2 * Definition for a binary tree node. 3 * struct TreeNode { 4 * int val; 5 * TreeNode *left; 6 * TreeNode *right; 7 * TreeNode() : val(0), left(nullptr), right(nullptr) {} 8 * TreeNode(int x) : val(x), left(nullptr), right(nullptr) {} 9 * TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {} 10 * }; 11 */ 12class Solution { 13public: 14 vector<int> treeQueries(TreeNode* root, vector<int>& queries) { 15 // Map to store the height of each node's subtree 16 unordered_map<TreeNode*, int> subtreeHeight; 17 18 // Calculate height of each subtree (post-order traversal) 19 // Height is defined as 1 + max(leftHeight, rightHeight) 20 function<int(TreeNode*)> calculateHeight = [&](TreeNode* node) -> int { 21 if (!node) { 22 return 0; 23 } 24 25 int leftHeight = calculateHeight(node->left); 26 int rightHeight = calculateHeight(node->right); 27 28 // Store the height of current node's subtree 29 subtreeHeight[node] = 1 + max(leftHeight, rightHeight); 30 return subtreeHeight[node]; 31 }; 32 33 // Calculate all subtree heights 34 calculateHeight(root); 35 36 // Array to store the maximum height of tree after removing each node 37 // Index represents node value, value represents max height after removal 38 vector<int> maxHeightAfterRemoval(subtreeHeight.size() + 1); 39 40 // DFS to calculate the maximum height when each node is removed 41 // depth: current depth in the tree (root starts at 0) 42 // maxHeightWithoutSubtree: max height achievable without current subtree 43 function<void(TreeNode*, int, int)> calculateMaxHeightAfterRemoval = 44 [&](TreeNode* node, int depth, int maxHeightWithoutSubtree) { 45 46 if (!node) { 47 return; 48 } 49 50 // Move to next depth level 51 ++depth; 52 53 // Store the maximum height if this node's subtree is removed 54 maxHeightAfterRemoval[node->val] = maxHeightWithoutSubtree; 55 56 // Process left child: max height is either the current max without subtree 57 // or the depth plus the height of the right subtree 58 calculateMaxHeightAfterRemoval( 59 node->left, 60 depth, 61 max(maxHeightWithoutSubtree, depth + subtreeHeight[node->right]) 62 ); 63 64 // Process right child: max height is either the current max without subtree 65 // or the depth plus the height of the left subtree 66 calculateMaxHeightAfterRemoval( 67 node->right, 68 depth, 69 max(maxHeightWithoutSubtree, depth + subtreeHeight[node->left]) 70 ); 71 }; 72 73 // Start DFS from root with initial depth -1 (so root has depth 0 after increment) 74 calculateMaxHeightAfterRemoval(root, -1, 0); 75 76 // Build answer array based on queries 77 vector<int> answer; 78 for (int nodeValue : queries) { 79 answer.emplace_back(maxHeightAfterRemoval[nodeValue]); 80 } 81 82 return answer; 83 } 84}; 85

1/** 2 * Definition for a binary tree node. 3 */ 4class TreeNode { 5 val: number; 6 left: TreeNode | null; 7 right: TreeNode | null; 8 constructor(val?: number, left?: TreeNode | null, right?: TreeNode | null) { 9 this.val = (val === undefined ? 0 : val); 10 this.left = (left === undefined ? null : left); 11 this.right = (right === undefined ? null : right); 12 } 13} 14 15function treeQueries(root: TreeNode | null, queries: number[]): number[] { 16 // Map to store the height of each node's subtree 17 const subtreeHeight = new Map<TreeNode, number>(); 18 19 /** 20 * Calculate height of each subtree using post-order traversal 21 * Height is defined as 1 + max(leftHeight, rightHeight) 22 * @param node - Current node being processed 23 * @returns Height of the subtree rooted at this node 24 */ 25 const calculateHeight = (node: TreeNode | null): number => { 26 if (!node) { 27 return 0; 28 } 29 30 const leftHeight = calculateHeight(node.left); 31 const rightHeight = calculateHeight(node.right); 32 33 // Store the height of current node's subtree 34 const height = 1 + Math.max(leftHeight, rightHeight); 35 subtreeHeight.set(node, height); 36 return height; 37 }; 38 39 // Calculate all subtree heights 40 calculateHeight(root); 41 42 // Array to store the maximum height of tree after removing each node 43 // Index represents node value, value represents max height after removal 44 const maxHeightAfterRemoval: number[] = new Array(100001).fill(0); 45 46 /** 47 * DFS to calculate the maximum height when each node is removed 48 * @param node - Current node being processed 49 * @param depth - Current depth in the tree (root starts at 0) 50 * @param maxHeightWithoutSubtree - Max height achievable without current subtree 51 */ 52 const calculateMaxHeightAfterRemoval = ( 53 node: TreeNode | null, 54 depth: number, 55 maxHeightWithoutSubtree: number 56 ): void => { 57 if (!node) { 58 return; 59 } 60 61 // Move to next depth level 62 depth++; 63 64 // Store the maximum height if this node's subtree is removed 65 maxHeightAfterRemoval[node.val] = maxHeightWithoutSubtree; 66 67 // Get heights of left and right subtrees (0 if null) 68 const leftSubtreeHeight = node.left ? subtreeHeight.get(node.left)! : 0; 69 const rightSubtreeHeight = node.right ? subtreeHeight.get(node.right)! : 0; 70 71 // Process left child: max height is either the current max without subtree 72 // or the depth plus the height of the right subtree 73 calculateMaxHeightAfterRemoval( 74 node.left, 75 depth, 76 Math.max(maxHeightWithoutSubtree, depth + rightSubtreeHeight) 77 ); 78 79 // Process right child: max height is either the current max without subtree 80 // or the depth plus the height of the left subtree 81 calculateMaxHeightAfterRemoval( 82 node.right, 83 depth, 84 Math.max(maxHeightWithoutSubtree, depth + leftSubtreeHeight) 85 ); 86 }; 87 88 // Start DFS from root with initial depth -1 (so root has depth 0 after increment) 89 calculateMaxHeightAfterRemoval(root, -1, 0); 90 91 // Build answer array based on queries 92 const answer: number[] = []; 93 for (const nodeValue of queries) { 94 answer.push(maxHeightAfterRemoval[nodeValue]); 95 } 96 97 return answer; 98} 99

Solution Approach

The solution uses two DFS traversals to efficiently precompute the answer for all possible queries.

First DFS - Calculate Subtree Heights:

The function f(root) performs a post-order traversal to calculate the height of each subtree:

Base case: If the node is None, return height 0
Recursively calculate heights of left (l) and right (r) subtrees
The current node's subtree height is 1 + max(l, r)
Store this in dictionary d where d[node] = height of subtree rooted at node

Second DFS - Calculate Heights After Removal:

The function dfs(root, depth, rest) precomputes what the tree height would be if each node's subtree were removed:

Parameters:

root: Current node being processed
depth: Current depth from the original root (starts at -1, incremented to 0 at root)
rest: Maximum height achievable without using the current node's subtree

The logic for each node:

Increment depth by 1 (current node's depth)
Store rest in res[root.val] - this is the tree height if we remove this node's subtree
Recursively process children:
- For left child: Pass max(rest, depth + d[root.right]) as the new rest
  - This considers either the existing rest or the path through the right sibling
- For right child: Pass max(rest, depth + d[root.left]) as the new rest
  - This considers either the existing rest or the path through the left sibling

Key Insight: When we're at a node and about to recurse into its left child, we calculate what the maximum height would be if the left subtree didn't exist. This would be either:

The rest value (paths not involving this node)
depth + d[root.right] (the path going through the right child instead)

Final Step:

After preprocessing, each query can be answered in O(1) time by looking up res[queries[i]].

Time Complexity: O(n) for preprocessing + O(m) for answering queries = O(n + m) Space Complexity: O(n) for storing the heights and results

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a small example to illustrate the solution approach.

Consider this binary tree:

 1  / \  2 3  /  4

Step 1: First DFS - Calculate Subtree Heights

We traverse post-order and calculate heights:

Node 4: No children, height = 0
Node 3: No children, height = 0
Node 2: Has left child (4), height = 1 + max(0, 0) = 1
Node 1: Has children (2,3), height = 1 + max(1, 0) = 2

Dictionary d after first DFS:

d[1] = 2 (entire tree height)
d[2] = 1 (subtree at node 2)
d[3] = 0 (subtree at node 3)
d[4] = 0 (leaf node)

Step 2: Second DFS - Calculate Heights After Removal

Starting at root with dfs(1, -1, 0):

At node 1 (depth = 0):

res[1] = 0 (the rest value - but node 1 is never queried as it's the root)
For left child (2): rest = max(0, 0 + d[3]) = max(0, 0) = 0
For right child (3): rest = max(0, 0 + d[2]) = max(0, 1) = 1

At node 2 (depth = 1, rest = 0):

res[2] = 0 (if we remove node 2's subtree, only node 1→3 path remains, height = 1)
Wait, let me recalculate: res[2] should be max(rest, depth + sibling) = max(0, 0 + 0) = 0
Actually res[2] = rest = 0 means the tree height without node 2's subtree
But we need to consider the path to node 3: the height would be 1 (root to node 3)

Let me recalculate more carefully:

At node 1 (depth = 0):

For left child (2): rest = max(0, 0 + 1 + d[3]) = max(0, 1) = 1 (The 1 represents going from root through right child)
For right child (3): rest = max(0, 0 + 1 + d[2]) = max(0, 2) = 2

At node 2 (depth = 1, rest = 1):

res[2] = 1 (removing node 2 leaves path 1→3 with height 1)
For left child (4): rest = max(1, 1 + 0) = 1

At node 3 (depth = 1, rest = 2):

res[3] = 2 (removing node 3 leaves path 1→2→4 with height 2)

At node 4 (depth = 2, rest = 1):

res[4] = 1 (removing node 4 leaves paths 1→2 and 1→3, max height is 1)

Step 3: Answer Queries

If queries = [2, 4, 3]:

Query 2: res[2] = 1 (tree becomes just 1→3)
Query 4: res[4] = 1 (tree becomes 1→2 and 1→3)
Query 3: res[3] = 2 (tree becomes 1→2→4)

Answer = [1, 1, 2]

This demonstrates how the algorithm precomputes all possible removal scenarios in just two tree traversals, allowing each query to be answered instantly.

Time and Space Complexity

Time Complexity: O(n + m)

The algorithm consists of two main phases:

The first phase computes heights using function f(root), which visits each node exactly once in a post-order traversal. This takes O(n) time where n is the number of nodes in the tree.
The second phase performs a pre-order traversal using dfs(root, depth, rest), visiting each node exactly once to compute the result for each node value. This also takes O(n) time.
Finally, constructing the answer array by iterating through the queries takes O(m) time where m is the number of queries.

Total time complexity: O(n) + O(n) + O(m) = O(n + m)

Space Complexity: O(n)

The space usage includes:

The dictionary d stores the height for each node in the tree, requiring O(n) space.
The res array has size len(d) + 1, which is O(n) space.
The recursive call stack for both f and dfs functions can go as deep as the height of the tree, which in the worst case (skewed tree) is O(n).
The output array for queries uses O(m) space, but this is typically not counted as auxiliary space since it's required for the output.

Total space complexity: O(n)

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Confusing Height vs Depth Definitions

One of the most common pitfalls is mixing up the definitions of height and depth, leading to off-by-one errors.

The Pitfall:

Height: Number of edges from a node to its furthest leaf (or sometimes nodes)
Depth: Number of edges from root to a node

In the given code, the calculate_height function actually returns the number of nodes (not edges) in the longest path from node to leaf, which can cause confusion.

Solution: Be consistent with your definition throughout the code. If the problem defines height as edges, adjust accordingly:

def calculate_height(node: Optional[TreeNode]) -> int:  if node is None:  return -1 # Return -1 for null nodes to count edges   left_height = calculate_height(node.left)  right_height = calculate_height(node.right)   # Height in terms of edges  node_heights[node] = 1 + max(left_height, right_height)  return node_heights[node]

2. Incorrect Handling of None Children

The Pitfall: When calculating max_depth_without_current for child nodes, accessing node_heights[None] returns 0 (due to defaultdict), but this might not align with your height definition.

For example, in this line:

max(max_depth_without_current, current_depth + node_heights[node.right])

If node.right is None, node_heights[None] returns 0, which might give incorrect results depending on your height definition.

Solution: Handle None children explicitly:

def compute_max_depth_without_subtree(node, current_depth, max_depth_without_current):  if node is None:  return   current_depth += 1  max_depth_after_removal[node.val] = max_depth_without_current   # Handle left child  if node.left:  right_subtree_contribution = 0 if node.right is None else node_heights[node.right]  compute_max_depth_without_subtree(  node.left,  current_depth,  max(max_depth_without_current, current_depth + right_subtree_contribution)  )   # Handle right child  if node.right:  left_subtree_contribution = 0 if node.left is None else node_heights[node.left]  compute_max_depth_without_subtree(  node.right,  current_depth,  max(max_depth_without_current, current_depth + left_subtree_contribution)  )

3. Array Size Miscalculation

The Pitfall: The code assumes node values are from 1 to n and creates an array of size len(node_heights) + 1. However, len(node_heights) includes the entry for None if any None nodes were processed.

Solution: Count actual nodes properly or use the given n value:

# Option 1: Count non-None nodes actual_nodes = sum(1 for k in node_heights.keys() if k is not None) max_depth_after_removal = [0] * (actual_nodes + 1)  # Option 2: Pass n as parameter if available def treeQueries(self, root: Optional[TreeNode], queries: List[int], n: int) -> List[int]:  max_depth_after_removal = [0] * (n + 1)

4. Initial Depth Value Confusion

The Pitfall: Starting current_depth at -1 in the initial call might seem counterintuitive and can lead to errors if not handled consistently.

Solution: Use clearer initial values and add comments:

# Start with depth 0 for root, which represents 0 edges from root to root compute_max_depth_without_subtree(root, 0, 0)  # Then adjust the increment logic inside the function def compute_max_depth_without_subtree(node, current_depth, max_depth_without_current):  if node is None:  return   max_depth_after_removal[node.val] = max_depth_without_current   # Depth for children is current_depth + 1  child_depth = current_depth + 1   # Process children with updated depth  compute_max_depth_without_subtree(  node.left,  child_depth,  max(max_depth_without_current, current_depth + node_heights[node.right])  )

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

How does quick sort divide the problem into subproblems?

Divide the array into two equal halves by index

Divide the array into two based on whether an element is smaller than an arbitrary value

Divide the array into a stray element and the rest of the array

Quick sort does not use divide and conquer

I don't know