leetcode hard : Height of Binary Tree After Subtree Removal Queries

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

 

dfs를 두번 돌려서

첫 dfs때 각 노드의 깊이와 height를 다 구해놓고

두번때 dfs때 해당 노드가 없을때의 maxheight값을 갱신하며 정답을 미리 pre-compute해놓으면 된다.

# Definition for a binary tree node.
# 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]:
        depth = {} 
        height = {}
        
        def find_depth_and_height(node, d):
            if not node: return 0
            depth[node.val] = d
            height[node.val] = 1 + max(find_depth_and_height(node.left, d+1), find_depth_and_height(node.right, d+1))
            return height[node.val]
        
        find_depth_and_height(root, 0)
        
        ans = {}
        def find_answer(node, maxheight):
            if not node:
                return 0
            
            ans[node.val] = maxheight
            lheight = depth[node.val] + (height[node.left.val] if node.left else 0)
            rheight = depth[node.val] + (height[node.right.val] if node.right else 0)
            find_answer(node.left, max(maxheight, rheight))
            find_answer(node.right, max(maxheight, lheight))
        find_answer(root, 0)    
            
        return [ans[i] for i in queries]

 

 

혼자 못풀어서 다음 솔루션으로부터 영감을 얻었다.

def treeQueries(self, root, queries, ans = {}) -> List[int]:
    @cache
    def height(r): return 1 + max(height(r.left), height(r.right)) if r else 0
    def dfs(r, depth, mx):
        if not r: return
        ans[r.val] = mx
        dfs(r.left, depth + 1, max(mx, depth + height(r.right)))
        dfs(r.right, depth + 1, max(mx, depth + height(r.left)))
    dfs(root, 0, 0)
    return [ans[v] for v in queries]