Height of Binary Tree After Subtree Removal Queries — Leetcode 2458

Original problem link: https://leetcode.com/problems/height-of-binary-tree-after-subtree-removal-queries/description/

In this problem, we’re given a binary tree with unique values from 1 to some n and an array of queries, where we need to return an array with the heights of the binary tree if queries[i] were to be removed from the tree.

The obvious brute force way to do it is O(n²) and is time limit exceeded. The tricky part is to ensure that we find the correct height for each removal since the removed node may or may not be on the longest path to a leaf. Thus, we calculate heights of subtrees, store them in a map and when precomputing the answers, keep the maximum height from previous levels.

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {

    Map<Integer, Integer> subtreeHeights;
    Map<Integer, Integer> answers;

    public int[] treeQueries(TreeNode root, int[] queries) {
        subtreeHeights = new HashMap<>();
        answers = new HashMap<>();

        findHeights(root);
        findAnswers(root.left, 0, root, true, 0);
        findAnswers(root.right, 0, root, false, 0);

        int[] result = new int[queries.length];
        for (int i = 0; i < queries.length; i++) {
            result[i] = answers.get(queries[i]);
        }

        return result;
    }

    private void findAnswers(TreeNode node, int parentLevel, TreeNode parent, boolean isLeftOfParent, int maxOtherHeight) {
        if (node == null) {
            return;
        }

        if (isLeftOfParent) {
            int otherSide = parent.right == null ? 0 : subtreeHeights.get(parent.right.val);
            answers.put(node.val, Math.max(parentLevel + otherSide, maxOtherHeight));
        } else {
            int otherSide = parent.left == null ? 0 : subtreeHeights.get(parent.left.val);
            answers.put(node.val, Math.max(parentLevel + otherSide, maxOtherHeight));
        }
        findAnswers(node.left, parentLevel + 1, node, true, answers.get(node.val));
        findAnswers(node.right, parentLevel + 1, node, false, answers.get(node.val));
    }

    private int findHeights(TreeNode node) {
        if (node == null) {
            return 0;
        }

        int result = 1 + Math.max(findHeights(node.left), findHeights(node.right));
        subtreeHeights.put(node.val, result);
        return result;
    }
}

This submission has 73 ms of runtime and beats 61% of submissions.