Reference: LeetCode
Difficulty: Medium

Problem

We are given the head node root of a binary tree, where additionally every node’s value is either a 0 or a 1. Return the same tree where every subtree (of the given tree) not containing a 1 has been removed.

Recall that the subtree of a node X is X, plus every node that is a descendant of X.

Note:

  • The binary tree will have at most 100 nodes.
  • The value of each node will only be 0 or 1.

Example: 

1
2
3
4
5
6
7
Example 1:
Input: [1,null,0,0,1]
Output: [1,null,0,null,1]
 
Explanation: 
Only the red nodes satisfy the property "every subtree not containing a 1".
The diagram on the right represents the answer.

1
2
3
Example 2:
Input: [1,0,1,0,0,0,1]
Output: [1,null,1,null,1]

Analysis

Methods:

  1. Recursion
    • Prune children of the tree recursively. The only decision at each node is whether to prune the left child or the right child or not.
    • containsOne(node): Whether the tree at this node contains 1, and it also prunes all subtrees not containing 1, e.g. node.left does not contain a 1, then we should prune it via node.left = null.
    • Also, the current node’s value needs to be checked.
    • Time: $O(N)$
    • Space: $O(h)$

Code

Recursion

Bad Code

Note:

  • This is the first version I wrote, but it turns out that containsOne(node) does a lot of repeated jobs.
  • To improve, consider the traversal order. We can put the line 1 after visiting left subtree and right subtree. Therefore, traversal order sometimes is very important in performance or in reducing repeated calculations.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
public TreeNode pruneTree(TreeNode root) {
  if (root == null) {
    return null;
  }
  if (!containsOne(root)) return null;
  root.left = pruneTree(root.left);
  root.right = pruneTree(root.right);
  return root;    
}

private boolean containsOne(TreeNode root) {
  if (root == null) {
    return false;
  }
  if (root.val == 1) return true; // line 1
  return containsOne(root.left) || containsOne(root.right); // line 2
}

Change pruneTree to:

1
return containsOne(root) ? root : null;

Change the line 1 & 2 to:

1
2
3
4
5
boolean l = containsOne(root.left); // do it first!
boolean r = containsOne(root.right);
if (l == false) root.left = null;   // pruning
if (r == false) root.right = null;
return (root.val == 1) || l || r;

I try using a HashMap to cache those calculated nodes, but the runtime does not improve and even get worse.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
public TreeNode pruneTree(TreeNode root) {
  HashMap<TreeNode, Boolean> map = new HashMap<>();
  return pruneTree(root, map);
}

private TreeNode pruneTree(TreeNode root, HashMap<TreeNode, Boolean> map) {
  if (root == null) {
    return null;
  }
  if (!containsOne(root, map)) return null;
  root.left = pruneTree(root.left, map);
  root.right = pruneTree(root.right, map);
  return root;    
}

private boolean containsOne(TreeNode root, HashMap<TreeNode, Boolean> map) {
  if (root == null) {
    return false;
  }
  if (map.containsKey(root)) return map.get(root); // cache found
  if (root.val == 1) { // set true
    map.put(root, true);
    return true;
  }
  boolean ret = containsOne(root.left, map) || containsOne(root.right, map);
  map.put(root, ret);
  return ret;
}

Clean Code

Note:

  • Can we do pruning in containsOne(node)?
  • It is a postorder traversal. Can we do it in preorder way, which means checking the current node first? No! If root.val == 1, we cannot just return true, since we need to check the children.
  • If root is null, we should return false. List all the base cases to see why it is the case.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
public TreeNode pruneTree(TreeNode root) {
  return containsOne(root) ? root : null;
}

private boolean containsOne(TreeNode root) {
  if (root == null) {
    return false;
  }

  boolean lRes = containsOne(root.left);
  boolean rRes = containsOne(root.right);
  if (lRes == false) root.left = null;    // pruning
  if (rRes == false) root.right = null;
  
  return (root.val == 1) || lRes || rRes; // if one of them contains 1, the tree should not be pruned.
}