2458. Height of Binary Tree After Subtree Removal Queries

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 ith 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 ith 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.

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
    int level[101010];
    unordered_map<int, vector<pair<int,int>>> mp;
    bool leaf(TreeNode* u) {
        return !u->left and !u->right;
    }
    int dfs(TreeNode* u, int dep, int par) {
        level[u->val] = dep;
        int res = dep;
        if(!leaf(u)) {
            if(u->left) res = max(res, dfs(u->left, dep + 1, u->val));
            if(u->right) res = max(res, dfs(u->right, dep + 1, u->val));
        }
        mp[dep].push_back({res,u->val});
        sort(rbegin(mp[dep]), rend(mp[dep]));
        if(mp[dep].size() > 2) mp[dep].pop_back();
        return res;
    }
public:
    vector<int> treeQueries(TreeNode* root, vector<int>& Q) {
        vector<int> res(Q.size());
        mp = unordered_map<int, vector<pair<int,int>>>();
        dfs(root, 0, -1);
        unordered_map<int, int> cache;
        for(int i = 0; i < Q.size(); i++) {
            if(cache.count(Q[i])) res[i] = cache[Q[i]];
            else {
                res[i] = level[Q[i]] - 1;
                int lvl = level[Q[i]];
                for(auto& [k,v] : mp[lvl]) {
                    if(v == Q[i]) continue;
                    res[i] = max(res[i], k);
                } 
            }
            cache[Q[i]] = res[i];
        }
        return res;
    }
};