LeetCode Solutions Blog

Height of Binary Tree After Subtree Removal Queries

Number: 2545

Difficulty: Hard

Paid? No

Companies: Google

Problem Description

Given a binary tree with unique node values and an array of queries, for each query you must remove the entire subtree rooted at the given node (note that the root is never removed). The query is independent – meaning after each removal the tree is restored to its initial state. Return an array where each entry is the height of the tree (defined as the number of edges on the longest path from the root to any node) after performing the subtree removal.

Key Insights

The removal of a subtree can only affect paths that go through the removed branch.
We can leverage a two-pass tree dynamic programming technique:
- The first DFS (bottom-up) computes the height (subtree-depth) of each node.
- The second DFS (top-down) propagates the best “alternative” height from ancestors (the value of the tree outside the current subtree).
For any query removal at node q, the new height of the tree becomes the maximum height that can be achieved without using any node in the subtree rooted at q. This is captured by the propagated value “up[q]” computed during the second DFS.

Space and Time Complexity

Time Complexity: O(n) for the two DFS traversals (n is the number of nodes).
Space Complexity: O(n) due to recursion stack and auxiliary arrays to store depths and parent-side values.

Solution

We solve the problem in two main steps:

Use a DFS (Depth-First Search) from the root to compute for each node the height of its subtree (i.e. the maximum distance from the node to a leaf in its subtree). This “subtree height” is stored for each node.
Use a second DFS to compute an “up” value for every node. The up value represents the maximum distance from the root to a node that does not lie in the current node’s subtree. For a child, this can be computed from its parent’s up value (increased by one) or, if the sibling branch offers a longer route, using that information.
With these two pieces of data, the answer for each query (removing the subtree rooted at node q) is simply up[q] because the maximum height of the tree post-removal will be determined by the best path from above q that avoids q's subtree.

Code Solutions

# Python solution with line-by-line comments

# 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: TreeNode, queries: list[int]) -> list[int]:
        # Dictionary to store each node's subtree height (number of edges to deepest leaf)
        subtreeHeight = {}
        # Dictionary to store the "up" value for each node.
        # "up[node]" represents the longest path from the root to some node that is not in node's subtree.
        up = {}
        # Dictionary to map node values to TreeNode for quick query lookup.
        nodes = {}
        
        # First DFS: compute subtree heights and record nodes in dictionary.
        def dfs1(node: TreeNode) -> int:
            if not node:
                return -1  # base value, so a leaf has height 0.
            nodes[node.val] = node
            leftHeight = dfs1(node.left)
            rightHeight = dfs1(node.right)
            # Height of current node is max height from left/right plus 1 edge.
            currHeight = max(leftHeight, rightHeight) + 1
            subtreeHeight[node.val] = currHeight
            return currHeight
        
        dfs1(root)
        
        # Second DFS: propagate "up" values.
        def dfs2(node: TreeNode, currUp: int):
            if not node:
                return
            # Store the up value for current node.
            up[node.val] = currUp
            # For computing child's up, determine the candidate coming from sibling.
            # For left child, the candidate from right subtree.
            leftCandidate = -1
            rightCandidate = -1
            if node.left and node.right:
                leftCandidate = subtreeHeight[node.right.val] + 1
                rightCandidate = subtreeHeight[node.left.val] + 1
            elif node.left:
                rightCandidate = 0  # No sibling value when only one child.
            elif node.right:
                leftCandidate = 0

            # Propagate to left child if exists.
            if node.left:
                leftUp = max(currUp + 1, leftCandidate + 1)
                dfs2(node.left, leftUp)
            # Propagate to right child if exists.
            if node.right:
                rightUp = max(currUp + 1, rightCandidate + 1)
                dfs2(node.right, rightUp)
        
        dfs2(root, 0)
        
        # For each query, answer is the up value for that node's removal.
        result = []
        for q in queries:
            # The answer is the maximum height we can achieve without using the subtree rooted at q.
            result.append(up[q])
        return result