Height of Binary Tree After Subtree Removal Queries

hard binary tree dfs precompute

Problem

You are given the root of a binary tree whose n nodes carry the distinct values 1…n, plus a list of queries. Each query names a non-root node: pretend to delete that node's entire subtree and report the height of what remains — the number of edges on the longest downward path from the root.

Queries are independent: the tree snaps back to its original shape before the next one, and answers must stay fast even with up to 10⁴ queries on 10⁵ nodes.

Inputroot = [5,8,9,2,1,3,7,4,6], queries = [3,2,4,8]

Output[3,2,3,2]

Removing node 2 also removes 4 and 6, leaving height 2; removing leaf 4 keeps 6 at depth 3, so the height stays 3.

root (level order, null for gaps)

queries (comma-separated node values)

def tree_queries(root, queries):
    best = {}                       # value -> height with that subtree removed
    max_h = 0                       # deepest level seen so far in this sweep

    def sweep(node, depth, left_first):
        nonlocal max_h
        if not node:
            return
        best[node.val] = max(best.get(node.val, 0), max_h)
        max_h = max(max_h, depth)
        kids = (node.left, node.right) if left_first else (node.right, node.left)
        for child in kids:
            sweep(child, depth + 1, left_first)

    sweep(root, 0, True)            # sweep 1: left-to-right preorder
    max_h = 0
    sweep(root, 0, False)           # sweep 2: right-to-left preorder
    return [best[q] for q in queries]

function treeQueries(root, queries) {
  const best = new Map();           // value -> height with that subtree removed
  let maxH = 0;                     // deepest level seen so far in this sweep

  function sweep(node, depth, leftFirst) {
    if (!node) return;
    best.set(node.val, Math.max(best.get(node.val) || 0, maxH));
    maxH = Math.max(maxH, depth);
    const kids = leftFirst ? [node.left, node.right] : [node.right, node.left];
    for (const child of kids) sweep(child, depth + 1, leftFirst);
  }

  sweep(root, 0, true);             // sweep 1: left-to-right preorder
  maxH = 0;
  sweep(root, 0, false);            // sweep 2: right-to-left preorder
  return queries.map(q => best.get(q));
}

class Solution {
    Map<Integer, Integer> best = new HashMap<>();
    int maxH = 0;

    public int[] treeQueries(TreeNode root, int[] queries) {
        sweep(root, 0, true);             // sweep 1: left-to-right preorder
        maxH = 0;
        sweep(root, 0, false);            // sweep 2: right-to-left preorder
        int[] ans = new int[queries.length];
        for (int i = 0; i < queries.length; i++) ans[i] = best.get(queries[i]);
        return ans;
    }

    void sweep(TreeNode node, int depth, boolean leftFirst) {
        if (node == null) return;
        best.merge(node.val, maxH, Math::max);
        maxH = Math.max(maxH, depth);
        TreeNode first = leftFirst ? node.left : node.right;
        TreeNode second = leftFirst ? node.right : node.left;
        sweep(first, depth + 1, leftFirst);
        sweep(second, depth + 1, leftFirst);
    }
}

class Solution {
public:
    unordered_map<int, int> best;
    int maxH = 0;

    vector<int> treeQueries(TreeNode* root, vector<int>& queries) {
        sweep(root, 0, true);             // sweep 1: left-to-right preorder
        maxH = 0;
        sweep(root, 0, false);            // sweep 2: right-to-left preorder
        vector<int> ans;
        for (int q : queries) ans.push_back(best[q]);
        return ans;
    }

    void sweep(TreeNode* node, int depth, bool leftFirst) {
        if (!node) return;
        best[node->val] = max(best[node->val], maxH);
        maxH = max(maxH, depth);
        TreeNode* first = leftFirst ? node->left : node->right;
        TreeNode* second = leftFirst ? node->right : node->left;
        sweep(first, depth + 1, leftFirst);
        sweep(second, depth + 1, leftFirst);
    }
};

Explanation

Recomputing the height after each removal costs O(n) per query — far too slow for 10⁴ queries. The trick: deleting the subtree of v keeps exactly the nodes outside that subtree, so the answer is just the deepest surviving level. We can precompute that for every node with two preorder sweeps.

Sweep 1 (left→right): when preorder first reaches v, every node visited so far — all ancestors plus everything to the left of v's subtree — survives the removal. So the running maximum depth maxH at that moment is a valid candidate, recorded as best[v] before maxH absorbs v's own depth.

Sweep 2 (right→left) mirrors the walk, visiting right children first. Now the nodes seen before v are its ancestors plus everything to the right of its subtree. Between the two sweeps every survivor is counted, so best[v] = max(both records) is exact — and the root always survives, so the value is well defined.

Walk the example root = [5,8,9,2,1,3,7,4,6]: for node 8, sweep 1 has only seen the root, recording 0; sweep 2 arrives after the whole right side, recording 2 (path 5→9→3). So removing 8 leaves height 2. For leaf 4, sweep 2 has already seen 6 at depth 3, so the height stays 3.

Each query q is then a single dictionary lookup best[q]. The default queries [3, 2, 4, 8] resolve instantly to [3, 2, 3, 2].

Both sweeps touch each node once, giving O(n) preprocessing and O(1) per query — O(n + q) total time with O(n) extra space for the table and recursion stack.

Time: O(n + q) Space: O(n)