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