2458. Height of Binary Tree After Subtree Removal Queries

Problem Statement

2458. Height of Binary Tree After Subtree Removal Queries

Hard

You are given the root of a binary tree with n nodes. Each node is assigned a unique value from 1 to n. You are also given an array queries of size m.

You have to perform m independent queries on the tree where in the i^th query you do the following:

Remove the subtree rooted at the node with the value queries[i] from the tree. It is guaranteed that queries[i] will not be equal to the value of the root.

Return an array answer of size m where answer[i] is the height of the tree after performing the i^th query.

Note:

The queries are independent, so the tree returns to its initial state after each query.
The height of a tree is the number of edges in the longest simple path from the root to some node in the tree.

Example 1:

Input: root = [1,3,4,2,null,6,5,null,null,null,null,null,7], queries = [4]
Output: [2]
Explanation: The diagram above shows the tree after removing the subtree rooted at node with value 4.
The height of the tree is 2 (The path 1 -> 3 -> 2).

Example 2:

Input: root = [5,8,9,2,1,3,7,4,6], queries = [3,2,4,8]
Output: [3,2,3,2]
Explanation: We have the following queries:
- Removing the subtree rooted at node with value 3. The height of the tree becomes 3 (The path 5 -> 8 -> 2 -> 4).
- Removing the subtree rooted at node with value 2. The height of the tree becomes 2 (The path 5 -> 8 -> 1).
- Removing the subtree rooted at node with value 4. The height of the tree becomes 3 (The path 5 -> 8 -> 2 -> 6).
- Removing the subtree rooted at node with value 8. The height of the tree becomes 2 (The path 5 -> 9 -> 3).

Constraints:

The number of nodes in the tree is n.
2 <= n <= 10⁵
1 <= Node.val <= n
All the values in the tree are unique.
m == queries.length
1 <= m <= min(n, 10⁴)
1 <= queries[i] <= n
queries[i] != root.val

Java — 2pass traversal

Source file

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 * int val;
 * TreeNode left;
 * TreeNode right;
 * TreeNode() {}
 * TreeNode(int val) { this.val = val; }
 * TreeNode(int val, TreeNode left, TreeNode right) {
 * this.val = val;
 * this.left = left;
 * this.right = right;
 * }
 * }
 */
class Solution {
    int maxHtYet = 0, idx = 0;

    public int[] treeQueries(TreeNode root, int[] queries) {
        Map<Integer, Integer> htAfterRm = new HashMap<>();
        traverseLtToRt(root, htAfterRm, 0);
        maxHtYet = 0;
        traverseRtToLt(root, htAfterRm, 0);
        int[] ans = new int[queries.length];
        for (int q : queries)
            ans[idx++] = htAfterRm.get(q);
        return ans;
    }

    private void traverseLtToRt(TreeNode root, Map<Integer, Integer> htAfterRm, int depth) {
        if (root == null)
            return;
        htAfterRm.put(root.val, maxHtYet);
        maxHtYet = Math.max(maxHtYet, depth);
        traverseLtToRt(root.left, htAfterRm, depth + 1);
        traverseLtToRt(root.right, htAfterRm, depth + 1);
    }

    private void traverseRtToLt(TreeNode root, Map<Integer, Integer> htAfterRm, int depth) {
        if (root == null)
            return;
        htAfterRm.put(root.val, Math.max(maxHtYet, htAfterRm.get(root.val)));
        maxHtYet = Math.max(maxHtYet, depth);
        traverseRtToLt(root.right, htAfterRm, depth + 1);
        traverseRtToLt(root.left, htAfterRm, depth + 1);
    }

}

Java

Source file

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 * int val;
 * TreeNode left;
 * TreeNode right;
 * TreeNode() {}
 * TreeNode(int val) { this.val = val; }
 * TreeNode(int val, TreeNode left, TreeNode right) {
 * this.val = val;
 * this.left = left;
 * this.right = right;
 * }
 * }
 */
class Solution {
    int maxHtYet = 0, idx = 0;

    public int[] treeQueries(TreeNode root, int[] queries) {
        Map<Integer, Integer> htAfterRm = new HashMap<>();
        traverseLtToRt(root, htAfterRm, 0);
        maxHtYet = 0;
        traverseRtToLt(root, htAfterRm, 0);
        int[] ans = new int[queries.length];
        for (int q : queries)
            ans[idx++] = htAfterRm.get(q);
        return ans;
    }

    private void traverseLtToRt(TreeNode root, Map<Integer, Integer> htAfterRm, int depth) {
        if (root == null)
            return;
        maxHtYet = Math.max(maxHtYet, depth);
        htAfterRm.put(root.val, maxHtYet);
        traverseLtToRt(root.left, htAfterRm, depth + 1);
        traverseLtToRt(root.right, htAfterRm, depth + 1);
    }

    private void traverseRtToLt(TreeNode root, Map<Integer, Integer> htAfterRm, int depth) {
        if (root == null)
            return;
        maxHtYet = Math.max(maxHtYet, depth);
        htAfterRm.put(root.val, Math.max(maxHtYet, htAfterRm.get(root.val)));
        traverseRtToLt(root.right, htAfterRm, depth + 1);
        traverseRtToLt(root.left, htAfterRm, depth + 1);
    }

}