1110.删点成林

48 阅读1分钟

题目:
给出二叉树的根节点 root,树上每个节点都有一个不同的值。

如果节点值在 to_delete 中出现,我们就把该节点从树上删去,最后得到一个森林(一些不相交的树构成的集合)。

返回森林中的每棵树。你可以按任意顺序组织答案。

解法:
删除指针P要使用**P=nil,指针指向的指针为nil

func delNodes(root *TreeNode, to_delete []int) []*TreeNode {
	// if root == nil {
	// 	return nil
	// }
	toDeleteMap := make(map[int]bool)
	for i := range to_delete {
		toDeleteMap[to_delete[i]] = true
	}
	ans := make([]*TreeNode, 0)
	var dfs func(root **TreeNode, isRoot bool)
	dfs = func(root **TreeNode, isRoot bool) {
		if *root == nil {
			return
		}

		_, delete := toDeleteMap[(*root).Val]

		if isRoot && !delete {
			ans = append(ans, *root)
		}

		if delete {
			dfs(&(*root).Left, true)
			dfs(&(*root).Right, true)

			destory(root)
		} else {
			dfs(&(*root).Left, false)
			dfs(&(*root).Right, false)
		}

	}
	dfs(&root, true)
	return ans
}

func destory(root **TreeNode) {
	if root == nil {
		return
	}
	*root = nil
}