1110. 删点成林

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

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

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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/delete-nodes-and-return-forest
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

class Solution {

    private Set<Integer> needDelete;

    private List<TreeNode> ans = new ArrayList<>();

    private void solve(TreeNode root, TreeNode parent, boolean canAdd) {
        if (root == null) {
            return;
        }

        if (!needDelete.contains(root.val)) {
            if (canAdd) {
                ans.add(root);
            }
            solve(root.left, root, false);
            solve(root.right, root, false);
        } else {
            if (parent != null) {
                if (parent.left == root) {
                    parent.left = null;
                } else {
                    parent.right = null;
                }
            }
            solve(root.left, root, true);
            solve(root.right, root, true);
        }
    }

    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        this.needDelete = new HashSet<Integer>() {{
            for (int x : to_delete) {
                add(x);
            }
        }};
        solve(root, null, true);
        return ans;
    }
}

class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;

    TreeNode() {
    }

    TreeNode(int val) {
        this.val = val;
    }

    TreeNode(int val, TreeNode left, TreeNode right) {
        this.val = val;
        this.left = left;
        this.right = right;
    }
}