1
\$\begingroup\$

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);
}
\$\endgroup\$
  • 3
    \$\begingroup\$ Welcome to Code Review. Is this code referring to the exercise binary tree pruning ? \$\endgroup\$ – dariosicily Feb 13 at 15:09
  • 2
    \$\begingroup\$ Please add at the challenge description to the question. \$\endgroup\$ – Mast Feb 14 at 6:34
  • \$\begingroup\$ @dariosicily yes \$\endgroup\$ – John Spring Feb 14 at 18:04
  • 1
    \$\begingroup\$ I've edited your post adding description of the exercise. \$\endgroup\$ – dariosicily Feb 15 at 7:35
2
\$\begingroup\$

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.

| improve this answer | |
\$\endgroup\$
0
\$\begingroup\$

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;
}
| improve this answer | |
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.