Leetcode — 814. Binary Tree Pruning

The link to the problem: https://leetcode.com/problems/binary-tree-pruning/

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.)

Example 1:
Input: [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.

photo credit to Leetcode

Thoughts:

Not a hard problem, but I just want to share the process of how I think up with this solution.

Basically, we are asked to remove all the “useless nodes”. Useless nodes will be the node that doesn’t contain any 1s either itself or in its child nodes, grandchild nodes, and etc.

When we traverse from top to bottom, it is hard to just immediately know we spot a useless node. We can only know the node might be useless if its value is not 1. We still need to traverse all the child nodes of the current node to make sure whether we should keep it or not.

Hence, it is better we traverse from bottom to top. Once we find a useless node, it means its parent node might in danger, i.e., its parent node might also be a useless node if the value of the parent node is 0 AND the parent node has another useless node. VICE VERSA —If we know a node is not a useless node, it means its parent node and grandparent nodes are safe! It’s just like the butterfly effect — whether a node is useless will affect all its parent nodes.

➡️Hence, the easiest way to solve this problem is using recursion so that we can check the leaf node first, and return null if we want to abandon the useless nodes.

Code Logic:

Since we are going to do the recursion, we need to decide our base case:

1. If the input node is null, then we can simply return null.
2. If the input node is a leaf node, then we can decide whether the current node is a useless node depends on its value. If its value is 0, then we are going to return null and pretend it doesn’t exist.

Ok, so let’s think about what are we going to do when the input node is not null or a leaf node.

Since we need to know whether the current node is useless depends on its two child nodes, we need to first check the status of them before deciding the status of the current node.

After checking the status of its two child nodes, we now need to check again if the current node contains 2 useless nodes. (i.e., see if the current node becomes leaf node since we ignore both its child nodes) Actually, this checking action is exactly the same as one of our base case. To make the whole code more simple and clean, we can actually move our “second base case” to here.

If the current node is not useless, then we just return the current node since we have also done pruning all its child nodes.

Below is the 0ms code that beats 100% —

public TreeNode pruneTree(TreeNode root) {
        if(root==null)
            return null;
        root.left = pruneTree(root.left);
        root.right = pruneTree(root.right);
        if(root.left==null && root.right==null)//if leaf node
            return root.val==0?null:root;
        return root;
    }