LeetCode/go/delete-nodes-and-return-forest.go

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
}