Leetcode Daily Question 26/10/2024 - 2458. Height of Binary Tree After Subtree Removal Queries

Problem Description:

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.

Solution: ►

from typing import Optional, List
class TreeNode:
  def __init__(self, val = 0, left = None, right = None):
    self.val = val
    self.left = left
    self.right = right


class Solution:
  def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
    removed_heights = defaultdict(int)
    self.curr_max = 0 
    
    def preorder(node, curr_height):
      if not node:
        return 
      
      removed_heights[node.val] = self.curr_max
      self.curr_max = max(self.curr_max, curr_height)
      
      preorder(node.left, curr_height + 1)
      preorder(node.right, curr_height + 1)
      
    
    def reversePreorder(node, curr_height):
      if not node:
        return
      
      removed_heights[node.val] = max(removed_heights[node.val], self.curr_max)
      self.curr_max = max(self.curr_max, curr_height)
      
      reversePreorder(node.right, curr_height + 1)
      reversePreorder(node.left, curr_height + 1)
      
    preorder(root, 0)
    self.curr_max = 0
    reversePreorder(root, 0)
    return [removed_heights[query] for query in queries]

Explanation:

We can solve this by traversing the tree and store the maximum height possible in memory if we remove the current node.

This way, we do not have to traverse the tree everytime we need to make a query, we can simply traverse the tree until we obtain all the information required.

We can use preorder traversal (root, left, right) first. The max height obtainable by removing the current note is the max between the height before reaching the current node, and the max height encountered.

However, just doing preorder is not enough. Consider the following example.

The preorder traversal only updates the information of height of right subtree nodes based on the height of the left subtree.

In the example above, if we remove node 2, we do not yet know that the right subtree has a deeper height than the left. So our output gives 3 instead of the expected 4.

To fix this, we have to do a right to left traversal to update the heights based on the information in the right subtree as well.

Let's define the preorder helper function first.

def preorder(node, curr_height):
      if not node:
        return 
      
      removed_heights[node.val] = self.curr_max
      self.curr_max = max(self.curr_max, curr_height)
      
      preorder(node.left, curr_height + 1)
      preorder(node.right, curr_height + 1)

Then let's do the "reverse" preorder (root, right, left).

def reversePreorder(node, curr_height):
      if not node:
        return
      
      removed_heights[node.val] = max(removed_heights[node.val], self.curr_max)
      self.curr_max = max(self.curr_max, curr_height)
      
      reversePreorder(node.right, curr_height + 1)
      reversePreorder(node.left, curr_height + 1)

To finish, we call preorder first, then reset the current max to 0 before we call reverse preorder to ensure the two traversal do not interfere with each other.

removed_heights = defaultdict(int)
self.curr_max = 0 # We store curr max globally, so it can be called by the helper functions.

# preorder(node, curr_height)

# reversePreorder(node, curr_height)

preorder(root, 0)
self.curr_max = 0
reversePreorder(root, 0)
return [removed_heights[query] for query in queries]

The time complexity is O(n). This is because we traverse the tree twice to obtain all the information needed to process the queries. O(2n) ≈ O(n)

The space complexity is O(n) as we may store as many queries as there are nodes.

The full solution can be found near the top of the page.

Back