Height of binary tree after subtree removal queries

Array Tree Depth-First Search Breadth-First Search Binary Tree

54.9% acceptance

Feb 27, 2026

1548

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 af ter performing the ith query. * Note: The queries are independent, so the tree returns to its initial state after e ach query. * The height of a tree is the number of edges in the longest simple path from t he root to some node in the tree. *

Solution

Rust

Time O(n)

Space O(n)

LeetCode

solution.rs

// Definition for a binary tree node.
// #[derive(Debug, PartialEq, Eq)]
// pub struct TreeNode {
//   pub val: i32,
//   pub left: Option<Rc<RefCell<TreeNode>>>,
//   pub right: Option<Rc<RefCell<TreeNode>>>,
// }
//
// impl TreeNode {
//   #[inline]
//   pub fn new(val: i32) -> Self {
//     TreeNode {
//       val,
//       left: None,
//       right: None
//     }
//   }
// }
use std::rc::Rc;
use std::cell::RefCell;
impl Solution {
  pub fn tree_queries(root: Option<Rc<RefCell<TreeNode>>>, queries: Vec<i32>) -> Vec<i32> {
    let n = 100001usize;
    let mut height = vec![0i32; n];
    let mut rest = vec![0i32; n];

    // Iterative post-order to compute heights
    let mut stack: Vec<(Rc<RefCell<TreeNode>>, bool)> = Vec::new();
    if let Some(r) = &root { stack.push((r.clone(), false)); }
    while let Some((node, processed)) = stack.pop() {
      if processed {
        let v = node.borrow();
        let lh = v.left.as_ref().map(|n| height[n.borrow().val as usize]).unwrap_or(-1);
        let rh = v.right.as_ref().map(|n| height[n.borrow().val as usize]).unwrap_or(-1);
        height[v.val as usize] = 1 + lh.max(rh);
      } else {
        stack.push((node.clone(), true));
        let v = node.borrow();
        if let Some(right) = &v.right { stack.push((right.clone(), false)); }
        if let Some(left) = &v.left { stack.push((left.clone(), false)); }
      }
    }

    // Iterative pre-order DFS to compute rest values
    // stack item: (node, depth, rest_val)
    let mut stack2: Vec<(Rc<RefCell<TreeNode>>, i32, i32)> = Vec::new();
    if let Some(r) = &root { stack2.push((r.clone(), 0, 0)); }
    while let Some((node, dep, rest_val)) = stack2.pop() {
      let v = node.borrow();
      rest[v.val as usize] = rest_val;
      let lh = v.left.as_ref().map(|n| height[n.borrow().val as usize]).unwrap_or(-1);
      let rh = v.right.as_ref().map(|n| height[n.borrow().val as usize]).unwrap_or(-1);
      let rest_left = rest_val.max(dep + if rh >= 0 { 1 + rh } else { 0 });
      let rest_right = rest_val.max(dep + if lh >= 0 { 1 + lh } else { 0 });
      if let Some(right) = &v.right { stack2.push((right.clone(), dep + 1, rest_right)); }
      if let Some(left) = &v.left { stack2.push((left.clone(), dep + 1, rest_left)); }
    }

    queries.iter().map(|&q| rest[q as usize]).collect()
  }
}