Binary tree pruning with BFS

Question

Below the text of the exercise:

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.

I have the following code and it does not look nice to me as too many if statements strolling around. How can I make this shorter and nice looking (more readable). Tree has nodes that contain 0 or 1 as value. I make the node null if it does not contain any node having 1 as value.

public TreeNode pruneTree(TreeNode root) {
 if (root == null || (root.left == null && root.right == null && root.val == 0)) return null;
 Queue<TreeNode> queue = new LinkedList<>();
 queue.offer(root);
 while(!queue.isEmpty()) {
 TreeNode node = queue.poll();
 if (node.left != null && !containsOne(node.left)) {
 node.left = null; 
 }
 if (node.right != null && !containsOne(node.right)) {
 node.right = null; 
 }
 if (node.left != null) {
 queue.offer(node.left);
 }
 if (node.right != null) {
 queue.offer(node.right);
 }
 }
 return root;
}
private boolean containsOne(TreeNode node) {
 if (node == null) return false;
 if (node.val == 1) return true;
 return containsOne(node.left) || containsOne(node.right);
}

• dariosicily: Welcome to Code Review. Is this code referring to the exercise binary tree pruning ?
• Mast: Please add at the challenge description to the question.
• John Spring: @dariosicily yes
• dariosicily: I've edited your post adding description of the exercise.

Answer 1

Note: I checked binary tree pruning and the structure of class TreeNode is not modifiable as I expected:

public class TreeNode {
 int val;
 TreeNode left;
 TreeNode right;
 TreeNode(int x) { val = x; }
}

It's possible to shorten your code with constraints of the class, your method containsOne is the following:

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

Because the third line will be executed when the condition node.val == 1 is false put directly this condition in the or expression in the third line:

private static boolean containsOne(TreeNode node) {
 if (node == null) return false;
 return node.val == 1 || containsOne(node.left) || containsOne(node.right);
}

About your method PruneTree you can shorten the following lines inside the method:

if (node.left != null && !containsOne(node.left)) {
 node.left = null; 
}
if (node.right != null && !containsOne(node.right)) {
 node.right = null; 
}
if (node.left != null) {
 queue.offer(node.left);
}
if (node.right != null) {
 queue.offer(node.right);
}

The code can be rewritten like below, two equal blocks and not so elegant to see but the original structure in the site cannot be modified, so I haven't thought about other alternatives:

if (node.left != null) {
 if (!containsOne(node.left)) { node.left = null; }
 else { queue.offer(node.left); }
}
if (node.right != null) {
 if (!containsOne(node.right)) { node.right = null; }
 else { queue.offer(node.right); }
}

I checked the code on the site passing all tests.

Answer 2

Unfortunately your attempt is far too complex and LeetCode.com doesn't notice that :(

The problem is that containsOne basically repeats the search for the subtrees that it already has run for in a previous iteration.

Example:

 A
 / \
 B C
 / \ 
D E

You start to search tree A which you do by searching subtrees B and C. In order to search B you search subtrees D and E. Then you move on to subtree B and repeat the search there, including repeating the searches on D and E. Next you move to subtree D and repeat the search there a third time, and so on.

Instead by using a recursive so-called post-order traversal where you basically start at the bottom (the leaves) of the tree, where you remove any leaves with the value 0, you get a much simpler solution:

public TreeNode pruneTree(TreeNode node) {
 if (node == null) {
 return null;
 }
 node.left = pruneTree(node.left);
 node.right = pruneTree(node.right);
 if (node.left == null && node.right == null && node.val == 0) {
 return null;
 }
 return node;
}