# Binary tree pruning with BFS

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.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);
}

• Welcome to Code Review. Is this code referring to the exercise binary tree pruning ? Feb 13 '20 at 15:09
– Mast
Feb 14 '20 at 6:34
• @dariosicily yes Feb 14 '20 at 18:04
• I've edited your post adding description of the exercise. Feb 15 '20 at 7:35

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.

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;
}