Skip to content

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<int> treeQueries(TreeNode* root, vector<int>& queries) {
    vector<int> ans;

    dfs(root, 0, 0);

    for (const int query : queries)
      ans.push_back(valToMaxHeight[query]);

    return ans;
  }

 private:
  // valToMaxHeight[val] := the maximum height without the node with `val`
  unordered_map<int, int> valToMaxHeight;
  // valToHeight[val] := the height of the node with `val`
  unordered_map<int, int> valToHeight;

  int height(TreeNode* root) {
    if (root == nullptr)
      return 0;
    if (const auto it = valToHeight.find(root->val); it != valToHeight.cend())
      return it->second;
    return valToHeight[root->val] =
               (1 + max(height(root->left), height(root->right)));
  }

  // maxHeight := the maximum height without the 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] := the maximum height without the node with `val`
  private Map<Integer, Integer> valToMaxHeight = new HashMap<>();
  // valToHeight[val] := the height of the node with `val`
  private Map<Integer, Integer> 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 := the maximum height without the 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: TreeNode | None, queries: list[int]) -> list[int]:
    @lru_cache(None)
    def height(root: TreeNode | None) -> int:
      if not root:
        return 0
      return 1 + max(height(root.left), height(root.right))

    # valToMaxHeight[val] := the maximum height without the node with `val`
    valToMaxHeight = {}

    # maxHeight := the maximum height without the current node `root`
    def dfs(root: TreeNode | None, 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]