LeetCode 1110 Delete Nodes And Return Forest (DFS 推荐)

Given the root of a binary tree, each node in the tree has a distinct value.

After deleting all nodes with a value in to_delete, we are left with a forest (a disjoint union of trees).

Return the roots of the trees in the remaining forest.  You may return the result in any order.

Example 1:

 

Input: root = [1,2,3,4,5,6,7], to_delete = [3,5]
Output: [[1,2,null,4],[6],[7]]

Constraints:

• The number of nodes in the given tree is at most 1000.
• Each node has a distinct value between 1 and 1000.
• to_delete.length <= 1000
• to_delete contains distinct values between 1 and 1000.

题目链接：https://leetcode.com/problems/delete-nodes-and-return-forest/

题目分析：比较适合面试的一道题，DFS时记录父亲结点以便于删除

1ms，时间击败100%

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    
    public void DFS(TreeNode cur, TreeNode fa, int[] deleteIDMap, List<TreeNode> ans) {
        if (deleteIDMap[cur.val] == 1) {
            if (cur.left != null) {
                if (deleteIDMap[cur.left.val] == 0) {
                    ans.add(cur.left);
                }
                DFS(cur.left, cur, deleteIDMap, ans);
            }
            if (cur.right != null) {
                if (deleteIDMap[cur.right.val] == 0) {
                    ans.add(cur.right);
                }
                DFS(cur.right, cur, deleteIDMap, ans);
            }
            if (fa == null) {
                return;
            }
            if (cur == fa.left) {
                fa.left = null;
            } else {
                fa.right = null;
            }
            return;
        }

        if (cur.left != null) {
            DFS(cur.left, cur, deleteIDMap, ans);
        }
        if (cur.right != null) {
            DFS(cur.right, cur, deleteIDMap, ans);
        }
    }
    
    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        int[] deleteIDMap = new int[1005];
        for (int i = 0; i < to_delete.length; i++) {
            deleteIDMap[to_delete[i]] = 1;
        }
        List<TreeNode> ans = new ArrayList<>();
        if (deleteIDMap[root.val] == 0) {
            ans.add(root);
        }
        DFS(root, null, deleteIDMap, ans);
        return ans;
    }
}