Leetcode每日一题练习 —— 1110. 删点成林（Delete Nodes And Return Forest）

题解

本题关键在于了解什么样的节点会成为森林中树的根节点, 根据描述可以得出结论: 如果父节点是一个要被删除的节点且子节点不被删除, 那么这个子节点将成为森林中一棵树的根节点. 使用dfs遍历树并根据这一结论将根节点加入最终结果中, 注意处理当父节点不被删除但其子节点被删除时将父节点的这个子节点设置为空. 根据我们的结论, 每次只要对父节点和其左右子节点做出判断即可, 这样就转换为了仅包含父节点和左右子节点的小问题, 就可以使用递归来遍历整棵树得出最终结果.

代码

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func delNodes(root *TreeNode, to_delete []int) []*TreeNode {
    to_delete_map := map[int]bool{}
    for _,value := range to_delete{
        to_delete_map[value] = true
    }
    result := []*TreeNode{}
    if !to_delete_map[root.Val]{
        result = append(result, root)
    }
    var dfs func(*TreeNode)
    dfs = func(root *TreeNode){
        if to_delete_map[root.Val]{
            if root.Left != nil{
                if !to_delete_map[root.Left.Val]{
                    result = append(result, root.Left)
                }
                dfs(root.Left)
            }
            if root.Right != nil{
                if !to_delete_map[root.Right.Val]{
                    result = append(result, root.Right)
                }
                dfs(root.Right)
            }
        }else{
            if root.Left != nil{
                dfs(root.Left)
                if to_delete_map[root.Left.Val]{
                    root.Left = nil
                }
            }
            if root.Right != nil{
                dfs(root.Right)
                if to_delete_map[root.Right.Val]{
                    root.Right = nil
                }
            }
        }  
    }
    dfs(root)
    return result
}