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 }