go: add «2458. Height of Binary Tree After Subtree Removal Queries»

URL:	https://leetcode.com/problems/height-of-binary-tree-after-subtree-removal-queries/
Signed-off-by: Matej Focko <me@mfocko.xyz>

go/height-of-binary-tree-after-subtree-removal-queries.go

@@ -0,0 +1,49 @@
+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
+}