Binary Tree Pruning

Problem Statement

Given the root of a binary tree, return the same tree where every subtree (of the given tree) not containing a 1 has been removed.

A subtree of a node node is node plus every node that is a descendant of node.

This problem could be found here :: https://leetcode.com/problems/binary-tree-pruning/

Example 1:

Input: root = [1,null,0,0,1]
Output: [1,null,0,null,1]
Explanation: 
Only the red nodes satisfy the property "every subtree not containing a 1".
The diagram on the right represents the answer.

Example 2:

Input: root = [1,0,1,0,0,0,1]
Output: [1,null,1,null,1]

Example 3:

Input: root = [1,1,0,1,1,0,1,0]
Output: [1,1,0,1,1,null,1]

Constraints:
The number of nodes in the tree is in the range [1, 200].
Node.val is either 0 or 1.

Solution

Let's look at how could we figure out the subtree's judgment whether it contains 1 or not. This could be done very naive technique of tree traversal, known DFS (depth first search).

One question persists how could the knowledge of subtree from depth be transferred from bottom to top?
Let's break down this, bit by bit. So, what kind of information needs to be sent, it is the understanding of whether 1 is present in the subtree or not (ie, boolean value needs to be returned).
Now, based on what we need to transfer the value, to above parent. For this we need to check

1. the current state the
2. the left node / subtree
3. the right node / subtree

then this, should be boil down to the pseudocode as below

current_node_has_one = node.value == 1;
left_node_has_one = check_for_one(node.left)
right_node_has_one = check_for_one(node.right)

now if left_node_has_one or right_node_has_one is true, then it's okay to keep it. If any one (or both) is false, then the corresponding subtree needs to be null, or in other words needs to be deleted.

if(right_node_has_one) 
then
    node.right = null // deleted the right node
end if

Now this needs to be transferred to the upper node. This data will be comprised of, the current sub tree had null or not. This means, current_node_has_one || left_node_has_one || right_node_has_one, and this logically needs to be in the end of the function. In the end, halt condition would be what if the node is null, then return false, as this does not has 1.

Based on this we can have the below solution.

Time Complexity: O(n)

Auxiliary Space Complexity: O(1)

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public TreeNode pruneTree(TreeNode root) {
        return fn(root)? root : null;
    }

    public boolean fn(TreeNode root) {
        if(root == null) return false;
        boolean hasOne = root.val == 1;
        boolean left = fn(root.left); if(!left) root.left = null;
        boolean right = fn(root.right); if(!right) root.right = null;
        return hasOne | left | right;
    }
}