Delete Nodes And Return Forest LeetCode solution in Java

1
class Solution {
2

3
  HashMap<Integer, List<TreeNode>> parent_val_child_nodes_map;
4
  HashMap<Integer, TreeNode> child_val_parent_node_map;
5

6
  public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
7

8
    // initialize map
9
    parent_val_child_nodes_map = new HashMap<>();
10
    child_val_parent_node_map = new HashMap<>();
11

12
    // fill map
13
    dfsFillMap(root);
14

15
    // traverse to_delete to find those that do not have parent after deleting it
16
    List<TreeNode> res = new ArrayList<>();
17

18
    // actually deleting nodes
19
    for (int delete_val : to_delete) {
20

21
      // if the node has parent
22
      if (child_val_parent_node_map.containsKey(delete_val)) {
23
        TreeNode parent_node = child_val_parent_node_map.get(delete_val);
24
        if (parent_node.left != null && parent_node.left.val == delete_val) {
25
          parent_node.left = null;
26
        }
27

28
        if (parent_node.right != null && parent_node.right.val == delete_val) {
29
          parent_node.right = null;
30
        }
31
      }
32
    }
33

34
    // add root to the list first because root has no parent
35
    // only if root.val is not in to_delete
36
    if (!Arrays.stream(to_delete).anyMatch(j -> j == root.val)) {
37
      res.add(root);
38
    }
39

40
    // add other nodes that do not have parent
41
    for (int delete_val : to_delete) {
42
      if (parent_val_child_nodes_map.containsKey(delete_val)) {
43
        for (int i = 0; i < parent_val_child_nodes_map.get(delete_val).size(); i++) {
44

45
          // make sure the add node is not in to_delete
46
          int child_node_val = parent_val_child_nodes_map.get(delete_val).get(i).val;
47
          if (!Arrays.stream(to_delete).anyMatch(j -> j == child_node_val)) {
48
            res.add(parent_val_child_nodes_map.get(delete_val).get(i));
49
          }
50
        }
51
      }
52
    }
53

54
    return res;
55
  }
56

57
  public void dfsFillMap(TreeNode root) {
58

59
    if (root == null) {
60
      return;
61
    }
62

63
    if (root.left != null) {
64
      parent_val_child_nodes_map.putIfAbsent(root.val, new ArrayList<>());
65
      parent_val_child_nodes_map.get(root.val).add(root.left);
66

67
      child_val_parent_node_map.putIfAbsent(root.left.val, root);
68
      dfsFillMap(root.left);
69
    }
70

71
    if (root.right != null) {
72
      parent_val_child_nodes_map.putIfAbsent(root.val, new ArrayList<>());
73
      parent_val_child_nodes_map.get(root.val).add(root.right);
74

75
      child_val_parent_node_map.putIfAbsent(root.right.val, root);
76
      dfsFillMap(root.right);
77
    }
78

79
    return;
80
  }
81
}