LC 6223. 移除子树后的二叉树高度
https://leetcode.cn/problems/height-of-binary-tree-after-subtree-removal-queries/
如果是每次现场计算，沿着修改节点向上溯源，那么最坏情况下面会出现O(n^2)的时间复杂度。正确的方式是：
- 首先进行一次预计算，得到节点左右子树的高度。
- 第二次遍历的时候，分别计算：
- 如果裁掉左子树，高度是多少？ `max(root.rh + d, h)`
- 如果裁掉右子树，高度是多少？ `max(root.lh + d, h)`
- 遍历queries去结果集合里面查询。
class Solution:
def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
index = {}
def prepare(root, p):
if root is None:
return 0
index[root.val] = root
root.parent = p
lh = prepare(root.left, root)
rh = prepare(root.right, root)
root.lh = lh
root.rh = rh
return max(lh, rh) + 1
prepare(root, None)
record = {}
def walk(root, h, d):
l = root.left
if l is not None:
h2 = max(root.rh + d, h)
record[l.val] = h2
walk(l, h2, d + 1)
r = root.right
if r is not None:
h2 = max(root.lh + d, h)
record[r.val] = h2
walk(r, h2, d + 1)
walk(root, 0, 0)
ans = []
for q in queries:
r = record[q]
ans.append(r)
return ans