# 2458. Height of Binary Tree After Subtree Removal Queries

• Time: $O(n)$
• Space: $O(n)$
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 class Solution { public: vector treeQueries(TreeNode* root, vector& queries) { vector ans; dfs(root, 0, 0); for (const int query : queries) ans.push_back(valToMaxHeight[query]); return ans; } private: // valToMaxHeight[val] := max height w/o node with val unordered_map valToMaxHeight; // valToHeight[val] := height of node with val unordered_map valToHeight; int height(TreeNode* root) { if (root == nullptr) return 0; if (const auto it = valToHeight.find(root->val); it != cend(valToHeight)) return it->second; return valToHeight[root->val] = (1 + max(height(root->left), height(root->right))); } // maxHeight := max height w/o current node root void dfs(TreeNode* root, int depth, int maxHeight) { if (root == nullptr) return; valToMaxHeight[root->val] = maxHeight; dfs(root->left, depth + 1, max(maxHeight, depth + height(root->right))); dfs(root->right, depth + 1, max(maxHeight, depth + height(root->left))); } }; 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 class Solution { public int[] treeQueries(TreeNode root, int[] queries) { int[] ans = new int[queries.length]; dfs(root, 0, 0); for (int i = 0; i < queries.length; ++i) ans[i] = valToMaxHeight.get(queries[i]); return ans; } // valToMaxHeight[val] := max height w/o node with val private Map valToMaxHeight = new HashMap<>(); // valToHeight[val] := height of node with val private Map valToHeight = new HashMap<>(); private int height(TreeNode root) { if (root == null) return 0; if (valToHeight.containsKey(root.val)) return valToHeight.get(root.val); final int h = 1 + Math.max(height(root.left), height(root.right)); valToHeight.put(root.val, h); return h; } // maxHeight := max height w/o current node root private void dfs(TreeNode root, int depth, int maxHeight) { if (root == null) return; valToMaxHeight.put(root.val, maxHeight); dfs(root.left, depth + 1, Math.max(maxHeight, depth + height(root.right))); dfs(root.right, depth + 1, Math.max(maxHeight, depth + height(root.left))); } } 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 class Solution: def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]: @lru_cache(None) def height(root: Optional[TreeNode]) -> int: if not root: return 0 return 1 + max(height(root.left), height(root.right)) # valToMaxHeight[val] := max height w/o node with val valToMaxHeight = {} # maxHeight := max height w/o current node root def dfs(root: Optional[TreeNode], depth: int, maxHeight: int) -> None: if not root: return valToMaxHeight[root.val] = maxHeight dfs(root.left, depth + 1, max(maxHeight, depth + height(root.right))) dfs(root.right, depth + 1, max(maxHeight, depth + height(root.left))) dfs(root, 0, 0) return [valToMaxHeight[query] for query in queries]