2458. Height of Binary Tree After Subtree Removal Queries

Problem Description

The problem presents a binary tree scenario where you have a tree with n nodes, each with a unique value assigned in the range from 1 to n. You are given an array called queries, which contains a list of m node values that you need to perform independent operations on.

For each value in the queries array, you are required to remove the subtree that is rooted at the node with the value of the query. It's important to note that none of the query values will ever be the value of the root node of the entire tree.

After performing each query -- which involves removing a subtree -- you need to calculate the height of the tree. The height is defined as the number of edges in the longest path from the root node to any leaf node.

The answer must be an array where the i-th element reflects the height of the tree after conducting the i-th query. Each query is independent, meaning that the state of the tree is reset back to its original state before the next query is performed.


To approach this solution, we consider two important tree-related metrics: the depth of a node and the height of a tree. These can be computed efficiently using a depth-first search (DFS) algorithm.

Here's the intuition broken down step-by-step:

  1. First, we want to perform a DFS traversal to find the 'depth' of each node's subtree. The 'depth' here refers to the maximum number of edges from the current node down to the leaf nodes in its subtree. This will be helpful later when we need to determine the height of the tree after removing subtrees.

  2. Once we compute the depth for each node, we execute another DFS traversal. During this traversal, for every node, we will maintain two pieces of information: the depth accumulated so far (from the root node), and the maximum value between the accumulated depth so far and the depth of the sibling's subtree. The latter is crucial because, when removing a node's subtree, the sibling's subtree might determine the new height of the entire tree.

  3. After these preparations, we perform the deletions as per the queries. Since each query is independent, we don't actually need to manipulate the tree after each query. The information gathered during the DFS traversals allows us to answer the queries directly.

  4. For each query, we access the pre-computed potential new height of the tree if the subtree rooted at that node was removed. We gathered this information during our second DFS traversal.

  5. Finally, we collect the new heights into an array in the same order as the queries, which gives us the required result.

This approach works because it efficiently pre-processes the tree to calculate depthen and to assess the impact of removing any given subtree on the height of the tree. By doing this pre-processing, answering each query becomes an O(1) operation, which is very efficient when dealing with a large number of queries.

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

Solution Approach

The solution involves two main parts: computation of depth of each subtree and a modified DFS to record the height of the tree post-removal of any node's subtree. Below is a step-by-step explanation:

  1. Depth Calculation using DFS: A helper function f is defined to calculate the depth of each node. It is a classic example of a post-order DFS because it computes the depth of child nodes first before the parent node. This function:

    • Terminates and returns 0 if the current node is None (base case for a leaf's child).
    • Recursively calculates the depth of the left and right subtrees of the current node.
    • Stores the calculated depth in a dictionary d where keys are the nodes themselves and values are the depths.
    • Returns the depth of the current node, which is 1 (for the current node) plus the maximum depth of its children to ensure we are considering the longest path.
  2. Preparation for Queries using DFS: Another function dfs is used for a depth-first traversal. Here's what happens during this traversal:

    • A call is made to dfs starting with the root node, an initial depth of -1, and a rest parameter initialized to 0. The rest parameter represents the maximum height of the tree if we removed the current node's subtree.
    • For each node, we compute and store its potential new height (after the hypothetical removal of a subtree) in an array res, indexed by the node's value.
    • When traversing left, the potential new height is the maximum between the already calculated rest and the depth of the right subtree plus the current depth.
    • The case is similar when traversing right: we consider the depth of the left subtree.
    • This traversal computes what the height of the tree would become if we were to remove the subtree rooted at each node.
  3. Handling Queries: With the pre-computed depth (d) and potential heights (res), the solution can process each query in queries by directly accessing the corresponding result from res. This is efficient as no further tree computations are needed when processing the queries.

1# Definition for a binary [tree](/problems/tree_intro) 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
7class Solution:
8    def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
9        """
10        The main solution function that processes tree queries and computes
11        the tree heights after removing subtrees as specified in each query.
12        """
14        # Initialize a defaultdict to store the depth of each node's subtree.
15        d = defaultdict(int)
17        # First DFS to calculate depths of all nodes' subtrees.
18        f(root)
20        # Array to store the potential new height if a subtree is removed.
21        res = [0] * (len(d) + 1)
23        # Second DFS to prepare for fast query handling.
24        dfs(root, -1, 0)
26        # Process each query to get the resulting [tree](/problems/tree_intro) heights.
27        return [res[v] for v in queries]

Overall, the solution leverages DFS for both depth computation and preparation for queries, with the aid of a dictionary to store depths and an array to store the potential new heights post-subtree removals. The power of this approach lies in its ability to answer queries in a very time-efficient manner (constant time per query) once the initial processing is done.

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

Depth first search is equivalent to which of the tree traversal order?

Example Walkthrough

Let's illustrate the solution approach with a simple example of a binary tree and a set of queries. Consider the following binary tree where numbers within the nodes represent their unique values:

1      1
2     / \
3    2   3
4   /   / \
5  4   5   6

Assume the following queries: [5, 3]. Now let's walk through how the solution would process this.

  1. Depth Calculation using DFS:

    • We start with our DFS from the root node 1. Node 1 has two children 2 and 3.
    • For Node 2, we go to the left child Node 4 which is a leaf node, so its depth is 0. This depth is stored in the dictionary d, so d[4]=0.
    • Since Node 2 is also a leaf node except for the already visited child, d[2] = d[4] + 1 = 1.
    • For Node 3, we calculate the depth for its children, Nodes 5 and 6, both being leaf nodes hence d[5]=0 and d[6]=0.
    • Node 3 then has a depth of d[3] = max(d[5], d[6]) + 1 = 1.
    • Finally, the depth of the root node 1 would be d[1] = max(d[2], d[3]) + 1 = 2.
  2. Preparation for Queries using DFS:

    • We conduct the DFS starting from the root again, where we calculate the res array.
    • At the root node 1, rest starts at 0. For its left child 2, the potential new height is max(0, d[3] + 1) which is 2.
      • We continue to Node 4, which is a leaf and has no impact on the res array at this point.
    • For the right child 3, the potential new height if 3 were removed is max(0, d[2] + 1) which is 2.
      • Similarly, its children 5 and 6 would both have res[5] and res[6] set to 2.
  3. Handling Queries:

    • When the query asks for [5, 3], we look up res[5] and res[3] in our res array.
    • res[5] is 2, and since removing Node 5 does not impact the height of the tree, the new height for the first query is 2.
    • res[3] is 2, and removing Node 3 along with its subtree also leaves the height unchanged since Node 2 is of the same height from the root, so the new height for the second query is 2.

Therefore, the final answer array after executing the queries [5, 3] on this tree would be [2, 2], suggesting that the height of the tree remains the same after the subtrees rooted at Node 5 or Node 3 are removed, for these particular queries.

This example demonstrates the computation of potential heights using DFS and how the solution enables constant-time response to queries by using pre-computed information.

Solution Implementation

1from collections import defaultdict
3# Definition for a binary tree node.
4class TreeNode:
5    def __init__(self, val=0, left=None, right=None):
6        self.val = val
7        self.left = left
8        self.right = right
10class Solution:
11    def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
12        # Helper function to calculate the depth of the tree
13        def calculate_depth(node):
14            if node is None:
15                return 0
16            # Recursively find the depth of the left and right subtrees
17            left_depth, right_depth = calculate_depth(node.left), calculate_depth(node.right)
18            # Store the depth of the current node
19            depth_dict[node] = 1 + max(left_depth, right_depth)
20            return depth_dict[node]
22        # Perform a Depth-First Search to find the highest visible value from each node
23        def dfs(node, current_depth, max_visible_value):
24            if node is None:
25                return
26            current_depth += 1
27            # Record the maximum visible value for the current node
28            results[node.val] = max_visible_value
29            # Explore the left and right subtrees
30            dfs(node.left, current_depth, max(max_visible_value, current_depth + depth_dict[node.right]))
31            dfs(node.right, current_depth, max(max_visible_value, current_depth + depth_dict[node.left]))
33        # Dictionary to hold the depth of each node
34        depth_dict = defaultdict(int)
35        # Calculate the depth of each node in the tree
36        calculate_depth(root)
37        # Initialize the results list with zeros for each value up to the number of nodes
38        results = [0] * (len(depth_dict) + 1)
39        # Start the DFS from the root node
40        dfs(root, -1, 0)
41        # Create the list of results for each query
42        return [results[value] for value in queries]
1import java.util.HashMap;
2import java.util.Map;
5 * Definition for a binary tree node.
6 */
7class TreeNode {
8    int val;
9    TreeNode left;
10    TreeNode right;
12    TreeNode() {}
14    TreeNode(int val) { this.val = val; }
16    TreeNode(int val, TreeNode left, TreeNode right) {
17        this.val = val;
18        this.left = left;
19        this.right = right;
20    }
23class Solution {
24    private Map<TreeNode, Integer> depthMap = new HashMap<>();
25    private int[] response;
27    // Main function to answer the queries based on the binary tree.
28    public int[] treeQueries(TreeNode root, int[] queries) {
29        calculateDepth(root);
30        response = new int[depthMap.size() + 1];
32        // Adding a base case to the map for null node.
33        depthMap.put(null, 0);
35        // Perform DFS to fill in the response array.
36        depthFirstSearch(root, -1, 0);
38        int queryCount = queries.length;
39        int[] answer = new int[queryCount];
40        for (int i = 0; i < queryCount; ++i) {
41            answer[i] = response[queries[i]];
42        }
43        return answer;
44    }
46    // Performs the DFS traversal to compute rest depth and updates the response.
47    private void depthFirstSearch(TreeNode node, int depth, int rest) {
48        if (node == null) {
49            return;
50        }
51        ++depth;
52        response[node.val] = rest;
53        depthFirstSearch(node.left, depth, Math.max(rest, depth + depthMap.get(node.right)));
54        depthFirstSearch(node.right, depth, Math.max(rest, depth + depthMap.get(node.left)));
55    }
57    // Helper function to compute the depth of each node.
58    private int calculateDepth(TreeNode node) {
59        if (node == null) {
60            return 0;
61        }
62        int leftDepth = calculateDepth(node.left);
63        int rightDepth = calculateDepth(node.right);
64        int maxDepth = 1 + Math.max(leftDepth, rightDepth);
65        depthMap.put(node, maxDepth);
66        return maxDepth;
67    }
1#include <vector>
2#include <unordered_map>
3#include <functional>
5// Definition for a binary tree node.
6struct TreeNode {
7    int val;
8    TreeNode *left;
9    TreeNode *right;
10    TreeNode() : val(0), left(nullptr), right(nullptr) {}
11    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
12    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
15class Solution {
17    // Method to process tree queries.
18    vector<int> treeQueries(TreeNode* root, vector<int>& queries) {
19        // Create a map to store the depth of each TreeNode.
20        unordered_map<TreeNode*, int> depthMap;
22        // Recursive lambda function to calculate depth of nodes.
23        // It returns the depth of the given node.
24        function<int(TreeNode*)> calculateDepth = [&](TreeNode* node) -> int {
25            if (!node) return 0;
26            int leftDepth = calculateDepth(node->left);
27            int rightDepth = calculateDepth(node->right);
28            depthMap[node] = 1 + max(leftDepth, rightDepth);
29            return depthMap[node];
30        };
31        // Call the recursive function to calculate depths.
32        calculateDepth(root);
34        // Create a results vector for the maximum rest (extra information you can gather while staying in that node)
35        // for each node value, plus 1 for the case where node value is 0.
36        vector<int> maximumRest(depthMap.size() + 1);
38        // Recursive lambda function for depth-first search to populate the rest values for each node.
39        function<void(TreeNode*, int, int)> dfs = [&](TreeNode* node, int depth, int rest) {
40            if (!node) return;
41            ++depth;
42            maximumRest[node->val] = rest;
44            // Visit left and right children with updated 'rest'
45            dfs(node->left, depth, max(rest, depth + depthMap[node->right]));
46            dfs(node->right, depth, max(rest, depth + depthMap[node->left]));
47        };
48        // Initialize DFS with root node, depth as -1, and rest as 0.
49        dfs(root, -1, 0);
51        // Answer vector to store results for the given queries
52        vector<int> answers;
54        // Loop over each query and fetch the corresponding rest value.
55        for (int value : queries) {
56            answers.push_back(maximumRest[value]);
57        }
58        return answers;
59    }
1// Node definition for a binary tree.
2class TreeNode {
3    val: number;
4    left: TreeNode | null;
5    right: TreeNode | null;
6    constructor(val = 0, left = null, right = null) {
7        this.val = val;
8        this.left = left;
9        this.right = right;
10    }
13// A map to store the depth of each TreeNode.
14const depthMap: Map<TreeNode, number> = new Map();
16// Recursive function to calculate the depth of nodes in the binary tree.
17// It returns the depth of the given node.
18const calculateDepth = (node: TreeNode | null): number => {
19    if (!node) return 0;
20    const leftDepth = calculateDepth(node.left);
21    const rightDepth = calculateDepth(node.right);
22    const depth = 1 + Math.max(leftDepth, rightDepth);
23    depthMap.set(node, depth);
24    return depth;
27// Vector to store the maximum rest (extra information that can be gathered while staying at that node)
28// for each node value. The array is initialized with a size that accommodates for node values starting at 0.
29let maximumRest: number[];
31// Recursive function for depth-first search to calculate the rest values for each node.
32const dfs = (node: TreeNode | null, depth: number, rest: number): void => {
33    if (!node) return;
34    depth++;
35    maximumRest[node.val] = rest;
37    // Visit the left and right children with updated 'rest'.
38    if (node.left) dfs(node.left, depth, Math.max(rest, depth + (depthMap.get(node.right) || 0)));
39    if (node.right) dfs(node.right, depth, Math.max(rest, depth + (depthMap.get(node.left) || 0)));
42// Function to process tree queries.
43// Given a root of a tree and an array of queries, it returns a vector of answers for each query.
44const treeQueries = (root: TreeNode, queries: number[]): number[] => {
45    // Calculate depths starting from the root.
46    calculateDepth(root);
48    // Initialize the maximumRest array with suitable size.
49    maximumRest = new Array(depthMap.size + 1).fill(0);
51    // Initialize DFS with the root node, a starting depth of -1, and an initial rest value of 0.
52    dfs(root, -1, 0);
54    // Answer array to store results corresponding to the given queries.
55    const answers: number[] = [];
57    // Process each query and populate answers with corresponding maximum rest values.
58    queries.forEach(value => {
59        answers.push(maximumRest[value]);
60    });
62    return answers;

Time and Space Complexity

The given Python code performs two depth-first searches (DFS) on a binary tree: the f function to compute the depths of all nodes and the dfs function to answer the queries.

Time Complexity

  1. The f function computes the depth (distance to the farthest leaf) for each node in the tree. It is a DFS that visits every node exactly once. The time complexity for this part is O(N), where N is the number of nodes in the tree.

  2. The dfs function is another DFS that also visits every node exactly once. It annotates each value with the maximum depth that can be reached either from its children or the rest of the tree (the "rest" value). The time complexity for this part is also O(N).

Since both parts are sequential, the overall time complexity remains O(N).

Space Complexity

  1. The space taken by the recursion stack during the DFS is O(H), where H is the height of the binary tree.

  2. The d defaultdict and the res list each store a value for every node in the tree. Together, they contribute O(N) space.

Considering the recursion stack and the data structures used, the worst-case space complexity is O(N), assuming the tree is skewed. In the best case, when the tree is balanced, the space complexity due to the recursion stack would be O(log N).

Combining these factors, the overall space complexity of the algorithm can be expressed as O(N) in the worst case and O(N + log N) in the best case (which simplifies to O(N) because O(N) dominates O(log N)).

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

Fast Track Your Learning with Our Quick Skills Quiz:

Given an array of 1,000,000 integers that is almost sorted, except for 2 pairs of integers. Which algorithm is fastest for sorting the array?

Recommended Readings

Got a question? Ask the Monster Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.

Tired of the LeetCode Grind?

Our structured approach teaches you the patterns behind problems, so you can confidently solve any challenge. Get started now to land your dream tech job.

Get Started