LeetCode 814. Binary Tree Pruning

原题链接在这里：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.

Example 2:
Input: [1,0,1,0,0,0,1]
Output: [1,null,1,null,1]

Example 3:
Input: [1,1,0,1,1,0,1,0]
Output: [1,1,0,1,1,null,1]

Note:

• The binary tree will have at most 100 nodes.
• The value of each node will only be 0 or 1.

题解：

Check left subtree is all zero, if yes, left child is null.

Check right subtree is all zero, if yes, right child is null.

If both left subtree and right subtree are all zero, and cur node's val is 0, then return true, which means current subtree is all zero.

Check back to root node, if all zero, then root is deleted.

Time Complexity: O(n).

Space: O(h).

AC Java:

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    public TreeNode pruneTree(TreeNode root) {
        if(root == null){
            return null;
        }
        
        boolean rootAllZero = isAllZero(root);
        if(rootAllZero){
            return null;
        }
        
        return root;
    }
    
    
    private boolean isAllZero(TreeNode root){
        if(root == null){
            return true;
        }
        
        boolean left = isAllZero(root.left);
        boolean right = isAllZero(root.right);
        if(left){
            root.left = null;
        }
        
        if(right){
            root.right = null;
        }
        
        if(left && right && root.val==0){
            return true;
        }
        
        return false;
    }
}