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 &amp;&amp; root.right == null &amp;&amp; root.val == 0)) return null;

    Queue&lt;TreeNode&gt; queue = new LinkedList&lt;&gt;();
    queue.offer(root);

    while(!queue.isEmpty()) {
        TreeNode node = queue.poll();

        if (node.left != null &amp;&amp; !containsOne(node.left)) {
            node.left = null; 
        }

        if (node.right != null &amp;&amp; !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);
}
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 &amp;&amp; !containsOne(node.left)) {
    node.left = null; 
}
if (node.right != null &amp;&amp; !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.