Height of Binary Tree After Subtree Removal Queries

class Solution {
public:
    vector<int> treeQueries(TreeNode* root, vector<int>& queries) {
        unordered_map<int,int> height, best;
        function<int(TreeNode*)> getHeight = [&](TreeNode* u) -> int {
            if (!u) return -1;
            int h = 1 + max(getHeight(u->left), getHeight(u->right));
            height[u->val] = h;
            return h;
        };
        getHeight(root);
        function<void(TreeNode*, int, int)> dfs = [&](TreeNode* u, int depth, int parentBest) {
            if (!u) return;
            best[u->val] = parentBest;
            int lh = u->left ? height[u->left->val] : -1;
            int rh = u->right ? height[u->right->val] : -1;
            dfs(u->left, depth + 1, max(parentBest, depth + 1 + rh));
            dfs(u->right, depth + 1, max(parentBest, depth + 1 + lh));
        };
        dfs(root, 0, 0);
        vector<int> ans;
        ans.reserve(queries.size());
        for (int q : queries) ans.push_back(best[q]);
        return ans;
    }
};

func treeQueries(root *TreeNode, queries []int) []int {
    height := map[int]int{}
    best := map[int]int{}
    maxI := func(a, b int) int {
        if a > b {
            return a
        }
        return b
    }
    var getHeight func(u *TreeNode) int
    getHeight = func(u *TreeNode) int {
        if u == nil {
            return -1
        }
        h := 1 + maxI(getHeight(u.Left), getHeight(u.Right))
        height[u.Val] = h
        return h
    }
    getHeight(root)
    var dfs func(u *TreeNode, depth, parentBest int)
    dfs = func(u *TreeNode, depth, parentBest int) {
        if u == nil {
            return
        }
        best[u.Val] = parentBest
        lh, rh := -1, -1
        if u.Left != nil {
            lh = height[u.Left.Val]
        }
        if u.Right != nil {
            rh = height[u.Right.Val]
        }
        dfs(u.Left, depth+1, maxI(parentBest, depth+1+rh))
        dfs(u.Right, depth+1, maxI(parentBest, depth+1+lh))
    }
    dfs(root, 0, 0)
    ans := make([]int, 0, len(queries))
    for _, q := range queries {
        ans = append(ans, best[q])
    }
    return ans
}

from typing import List, Optional

class Solution:
    def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
        height: dict = {}
        best: dict = {}

        def get_height(u: Optional[TreeNode]) -> int:
            if not u:
                return -1
            h = 1 + max(get_height(u.left), get_height(u.right))
            height[u.val] = h
            return h

        get_height(root)

        def dfs(u: Optional[TreeNode], depth: int, parent_best: int) -> None:
            if not u:
                return
            best[u.val] = parent_best
            lh = height[u.left.val] if u.left else -1
            rh = height[u.right.val] if u.right else -1
            dfs(u.left, depth + 1, max(parent_best, depth + 1 + rh))
            dfs(u.right, depth + 1, max(parent_best, depth + 1 + lh))

        dfs(root, 0, 0)
        return [best[q] for q in queries]

function treeQueries(root, queries) {
    const height = new Map();
    const best = new Map();
    function getHeight(u) {
        if (!u) return -1;
        const h = 1 + Math.max(getHeight(u.left), getHeight(u.right));
        height.set(u.val, h);
        return h;
    }
    getHeight(root);
    function dfs(u, depth, parentBest) {
        if (!u) return;
        best.set(u.val, parentBest);
        const lh = u.left ? height.get(u.left.val) : -1;
        const rh = u.right ? height.get(u.right.val) : -1;
        dfs(u.left, depth + 1, Math.max(parentBest, depth + 1 + rh));
        dfs(u.right, depth + 1, Math.max(parentBest, depth + 1 + lh));
    }
    dfs(root, 0, 0);
    return queries.map((q) => best.get(q));
}

class Solution {
    private Map<Integer, Integer> height = new HashMap<>();
    private Map<Integer, Integer> best = new HashMap<>();

    public int[] treeQueries(TreeNode root, int[] queries) {
        getHeight(root);
        dfs(root, 0, 0);
        int[] ans = new int[queries.length];
        for (int i = 0; i < queries.length; i++) ans[i] = best.get(queries[i]);
        return ans;
    }

    private int getHeight(TreeNode u) {
        if (u == null) return -1;
        int h = 1 + Math.max(getHeight(u.left), getHeight(u.right));
        height.put(u.val, h);
        return h;
    }

    private void dfs(TreeNode u, int depth, int parentBest) {
        if (u == null) return;
        best.put(u.val, parentBest);
        int lh = u.left != null ? height.get(u.left.val) : -1;
        int rh = u.right != null ? height.get(u.right.val) : -1;
        dfs(u.left, depth + 1, Math.max(parentBest, depth + 1 + rh));
        dfs(u.right, depth + 1, Math.max(parentBest, depth + 1 + lh));
    }
}

function treeQueries(root: TreeNode | null, queries: number[]): number[] {
    const height = new Map<number, number>();
    const best = new Map<number, number>();
    const getHeight = (u: TreeNode | null): number => {
        if (!u) return -1;
        const h = 1 + Math.max(getHeight(u.left), getHeight(u.right));
        height.set(u.val, h);
        return h;
    };
    getHeight(root);
    const dfs = (u: TreeNode | null, depth: number, parentBest: number): void => {
        if (!u) return;
        best.set(u.val, parentBest);
        const lh = u.left ? height.get(u.left.val)! : -1;
        const rh = u.right ? height.get(u.right.val)! : -1;
        dfs(u.left, depth + 1, Math.max(parentBest, depth + 1 + rh));
        dfs(u.right, depth + 1, Math.max(parentBest, depth + 1 + lh));
    };
    dfs(root, 0, 0);
    return queries.map((q) => best.get(q)!);
}