LeetCode/go/height-of-binary-tree-after-subtree-removal-queries.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
}