我寫完後才發現
每個query[i]都去跑一次dfs也可以過
那這題根本是easy阿
到底三小
2458. Height of Binary Tree After Subtree Removal Queries
給一個二元樹的root有n個node
每個node都有一個唯一的值從1~n
給一個array : query
ans[i]為把node query[i]從二元樹刪掉後的高度
請回傳ans[i]
思路:
建立一個int[][3]{}的array
紀錄每一層(level)的最大高度、該層中高度為最大高度的node的個數、第二大的高度
建立一個hash table
紀錄每個node的level和高度
接著就跑一次dfs去把上面的array和hash table填好
接著按照query去跑
(1)
如果query[i]的高度 < 這一層的最大高度
ans[i] = 這一層的最大高度
(2)
如果query[i]的高度 = 這一層的最大高度 且 這一層高度為最大高度的node個數 > 1
ans[i] = 這一層的最大高度
(3)
如果query[i]的高度 = 這一層的最大高度 且 這一層高度為最大高度的node個數 = 1
ans[i] = 這一層第二高的高度
照上面那樣遍歷完所有query[i]
最後回傳ans就好
golang code :
/**
* Definition for a binary tree node.
* type TreeNode struct {
* Val int
* Left *TreeNode
* Right *TreeNode
* }
*/
func treeQueries(root *TreeNode, queries []int) []int {
max_depth_cnt := [][3]int{} // max_depth max_depth_cnt second_depth
node_level_depth := make(map[int][2]int)
ans := make([]int, len(queries))
dfs(0, root, &max_depth_cnt, &node_level_depth)
for key, val := range queries {
level := node_level_depth[val][0]
depth := node_level_depth[val][1]
if max_depth_cnt[level][0] > depth {
ans[key] = max_depth_cnt[level][0]
} else if max_depth_cnt[level][0] == depth {
if max_depth_cnt[level][1] == 1 {
ans[key] = max(level-1, max_depth_cnt[level][2])
} else {
ans[key] = max_depth_cnt[level][0]
}
}
}
return ans
}
func dfs(level int, node *TreeNode, max_depth_cnt *[][3]int, node_level_depth
*map[int][2]int) int {
if node == nil {
return level - 1
}
if len(*max_depth_cnt) <= level {
*max_depth_cnt = append(*max_depth_cnt, [3]int{})
}
cur_node_depth_left := dfs(level+1, node.Left, max_depth_cnt, node_level_
depth)
cur_node_depth_right := dfs(level+1, node.Right, max_depth_cnt, node_level_
depth)
cur_node_depth := max(cur_node_depth_left, cur_node_depth_right)
if cur_node_depth > (*max_depth_cnt)[level][0] {
(*max_depth_cnt)[level][2] = (*max_depth_cnt)[level][0]
(*max_depth_cnt)[level][0] = cur_node_depth
(*max_depth_cnt)[level][1] = 1
} else if cur_node_depth == (*max_depth_cnt)[level][0] {
(*max_depth_cnt)[level][1]++
} else if cur_node_depth > (*max_depth_cnt)[level][2] {
(*max_depth_cnt)[level][2] = cur_node_depth
}
(*node_level_depth)[node.Val] = [2]int{level, cur_node_depth}
return cur_node_depth
}