50 lines
1 KiB
Go
50 lines
1 KiB
Go
|
package main
|
||
|
|
|
||
|
|
func treeQueries(root *TreeNode, queries []int) []int {
|
||
|
|
depths := make(map[int]int)
|
||
|
|
subtreeHeights := make(map[int]int)
|
||
|
|
levelHeights := make(map[int][]int)
|
||
|
|
|
||
|
|
var dfs func(*TreeNode, int) int
|
||
|
|
dfs = func(node *TreeNode, level int) int {
|
||
|
|
if node == nil {
|
||
|
|
return 0
|
||
|
|
}
|
||
|
|
|
||
|
|
depths[node.Val] = level
|
||
|
|
|
||
|
|
left, right := dfs(node.Left, level+1), dfs(node.Right, level+1)
|
||
|
|
height := 1 + max(left, right)
|
||
|
|
subtreeHeights[node.Val] = height
|
||
|
|
|
||
|
|
maxHeights, found := levelHeights[level]
|
||
|
|
if !found {
|
||
|
|
levelHeights[level] = []int{height, 0}
|
||
|
|
} else if height > maxHeights[0] {
|
||
|
|
maxHeights[0], maxHeights[1] = height, maxHeights[0]
|
||
|
|
} else if height > maxHeights[1] {
|
||
|
|
maxHeights[1] = height
|
||
|
|
}
|
||
|
|
|
||
|
|
return height
|
||
|
|
}
|
||
|
|
dfs(root, 0)
|
||
|
|
|
||
|
|
results := make([]int, len(queries))
|
||
|
|
|
||
|
|
for i, node := range queries {
|
||
|
|
level := depths[node]
|
||
|
|
|
||
|
|
extremas, _ := levelHeights[level]
|
||
|
|
|
||
|
|
results[i] = level - 1
|
||
|
|
if subtreeHeights[node] == extremas[0] {
|
||
|
|
results[i] += extremas[1]
|
||
|
|
} else {
|
||
|
|
results[i] += extremas[0]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return results
|
||
|
|
}
|