2458.js

/** ------------------------------------------------------------------------------
 * 
 * 2458. Height of Binary Tree After Subtree Removal Queries
 * Topics: Binary Tree
 * https://leetcode.com/problems/height-of-binary-tree-after-subtree-removal-queries/?envType=daily-question&envId=2024-10-26
 * 
 ------------------------------------------------------------------------------ */
var treeQueries = function (root, queries) {
  if (!root) return []

  let ans = {}
  let heights = {}

  function maxHeight(node) {
    if (!node) return 0

    if (heights[node.val]) return heights[node.val]

    heights[node.val] = 1 + Math.max(maxHeight(node.left), maxHeight(node.right))

    return heights[node.val]
  }

  function dfs(node, depth, maxDepthWithoutCurrentNode) {
    if (!node) return 0

    ans[node.val] = maxDepthWithoutCurrentNode

    dfs(node.left, depth + 1, Math.max(maxDepthWithoutCurrentNode, depth + maxHeight(node.right)))
    dfs(node.right, depth + 1, Math.max(maxDepthWithoutCurrentNode, depth + maxHeight(node.left)))
  }

  dfs(root, 0, 0)

  return queries.map((q) => ans[q])
}