diff --git a/go/delete-nodes-and-return-forest.go b/go/delete-nodes-and-return-forest.go new file mode 100644 index 0000000..fec92ed --- /dev/null +++ b/go/delete-nodes-and-return-forest.go @@ -0,0 +1,38 @@ +package main + +import "slices" + +func delNodes(root *TreeNode, to_delete []int) []*TreeNode { + var forest []*TreeNode + + var delNodesRec func(*TreeNode) *TreeNode + delNodesRec = func(node *TreeNode) *TreeNode { + if node == nil { + return nil + } + + node.Left = delNodesRec(node.Left) + node.Right = delNodesRec(node.Right) + + if slices.Contains(to_delete, node.Val) { + if node.Left != nil { + forest = append(forest, node.Left) + } + + if node.Right != nil { + forest = append(forest, node.Right) + } + + return nil + } + + return node + } + + root = delNodesRec(root) + if root != nil { + forest = append(forest, root) + } + + return forest +}