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

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