From 8991ba7ff170d58bda774181fbb634a114a22846 Mon Sep 17 00:00:00 2001
From: Matej Focko <me@mfocko.xyz>
Date: Sun, 27 Oct 2024 00:08:01 +0200
Subject: [PATCH] =?UTF-8?q?go:=20add=20=C2=AB2458.=20Height=20of=20Binary?=
 =?UTF-8?q?=20Tree=20After=20Subtree=20Removal=20Queries=C2=BB?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

URL:	https://leetcode.com/problems/height-of-binary-tree-after-subtree-removal-queries/
Signed-off-by: Matej Focko <me@mfocko.xyz>
---
 ...nary-tree-after-subtree-removal-queries.go | 49 +++++++++++++++++++
 1 file changed, 49 insertions(+)
 create mode 100644 go/height-of-binary-tree-after-subtree-removal-queries.go

diff --git a/go/height-of-binary-tree-after-subtree-removal-queries.go b/go/height-of-binary-tree-after-subtree-removal-queries.go
new file mode 100644
index 0000000..73c8a8c
--- /dev/null
+++ b/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
+}