Height of Binary Tree After Subtree Removal Queries Problem

Input: root = [1,3,4,2,null,6,5,null,null,null,null,null,7], queries = [4]
Output: [2]
Explanation: The diagram above shows the tree after removing the subtree rooted at node with value 4.
The height of the tree is 2 (The path 1 -> 3 -> 2).

Input: root = [5,8,9,2,1,3,7,4,6], queries = [3,2,4,8]
Output: [3,2,3,2]
Explanation: We have the following queries:
- Removing the subtree rooted at node with value 3. The height of the tree becomes 3 (The path 5 -> 8 -> 2 -> 4).
- Removing the subtree rooted at node with value 2. The height of the tree becomes 2 (The path 5 -> 8 -> 1).
- Removing the subtree rooted at node with value 4. The height of the tree becomes 3 (The path 5 -> 8 -> 2 -> 6).
- Removing the subtree rooted at node with value 8. The height of the tree becomes 2 (The path 5 -> 9 -> 3).

public class Solution {
    public int[] treeQueries(TreeNode root, int[] queries) {
        int originalHeight = getHeight(root);
        int[] result = new int[queries.length];
        
        for (int i = 0; i < queries.length; i++) {
            result[i] = dfs(root, queries[i]);
        }
        
        return result;
    }
    public int dfs(TreeNode node, int block) {
        if (node == null || node.val == block) {
	        return -1;
	    }
        return 1 + Math.max(dfs(node.left, block), dfs(node.right, block));
    }    
    public int getHeight(TreeNode node) {
        if (node == null) return -1;
        return 1 + Math.max(getHeight(node.left), getHeight(node.right));
    }
}

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Solution:
    def getHeight(self, node):
        if not node:
            return -1
        return 1 + max(self.getHeight(node.left), self.getHeight(node.right))
    
    def dfs(self, node, block):
        if not node or node.val == block:
            return -1
        return 1 + max(self.dfs(node.left, block), self.dfs(node.right, block))
    
    def treeQueries(self, root, queries):
        originalHeight = self.getHeight(root)
        result = []
        
        for q in queries:
            result.append(self.dfs(root, q))
        
        return result

public class Solution {
    // Maps to store the depth and height of each node
    private Map<Integer, Integer> nodeDepths = new HashMap<>();
    private Map<Integer, Integer> nodeHeights = new HashMap<>();

    // Main function to handle the tree queries
    public int[] treeQueries(TreeNode root, int[] queries) {
        // Compute depth and height for each node
        computeDepthAndHeight(root, 0);

        // Map to store nodes according to their depth and heights
        Map<Integer, List<int[]>> peers = new HashMap<>();
        for (Map.Entry<Integer, Integer> entry : nodeDepths.entrySet()) {
            int val = entry.getKey();
            int depth = entry.getValue();

            // Initialize list for this depth if not already present
            peers.putIfAbsent(depth, new ArrayList<>());
            List<int[]> peerList = peers.get(depth);

            // Add the node's height and value to the list
            peerList.add(new int[]{nodeHeights.get(val), val});

            // Sort the list to keep nodes with the highest heights first
            peerList.sort((a, b) -> b[0] - a[0]);
            // Keep only the top 2 heights for this depth
            if (peerList.size() > 2) {
                peerList.remove(peerList.size() - 1);
            }
        }

        // Array to store the results of each query
        int[] result = new int[queries.length];
        // Process each query
        for (int i = 0; i < queries.length; i++) {
            int q = queries[i];
            int depth = nodeDepths.get(q);

            List<int[]> peerList = peers.get(depth);
            // Determine the longest path after removing the node
            if (peerList.size() == 1) {  // No peer, path length equals depth - 1
                result[i] = depth - 1;
            } else if (peerList.get(0)[1] == q) {  // The removed node has the largest height
                result[i] = peerList.get(1)[0] + depth;
            } else {  // Find the node with the largest height
                result[i] = peerList.get(0)[0] + depth;
            }
        }

        return result;
    }

    // DFS to compute depth and height
    private int computeDepthAndHeight(TreeNode node, int depth) {
        if (node == null) return -1;
        
        // Store the depth of the current node
        nodeDepths.put(node.val, depth);

        // Recursively compute the heights of left and right children
        int leftHeight = computeDepthAndHeight(node.left, depth + 1);
        int rightHeight = computeDepthAndHeight(node.right, depth + 1);
        
        // Compute the current node's height
        int height = 1 + Math.max(leftHeight, rightHeight);

        // Store the height of the current node
        nodeHeights.put(node.val, height);
        return height;
    }
}

class Solution:
    def treeQueries(self, R: Optional[TreeNode], Q: List[int]) -> List[int]:
        Depth = collections.defaultdict(int)
        Height = collections.defaultdict(int)

        def dfs(node, depth):
            if not node:
                return -1
            Depth[node.val] = depth
            cur = max(dfs(node.left, depth + 1), dfs(node.right, depth + 1)) + 1
            Height[node.val] = cur
            return cur

        dfs(R, 0)

        peers = collections.defaultdict(list)  # Group nodes according to their depth, keep the top 2 heights.
        for val, depth in Depth.items():
            peers[depth].append((-Height[val], val))  # Use negative heights for sorting in descending order
            peers[depth].sort()
            if len(peers[depth]) > 2:
                peers[depth].pop()

        ans = []
        for q in Q:
            depth = Depth[q]
            if len(peers[depth]) == 1:  # No peer, path length equals depth - 1.
                ans.append(depth - 1)
            elif peers[depth][0][1] == q:  # The removed node has the largest height, look for the node with 2nd largest height.
                ans.append(-peers[depth][1][0] + depth)
            else:  # Look for the node with the largest height.
                ans.append(-peers[depth][0][0] + depth)
        return ans