[LintCode/LeetCode] Binary Tree Pruning

Problem

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

Example 1:
Input: [1,#,0,0,1]
Output: [1,#,0,#,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,#,1,#,1]

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

Solution

public class Solution {
    /**
     * @param root: the root
     * @return: the same tree where every subtree (of the given tree) not containing a 1 has been removed
     */
    public TreeNode pruneTree(TreeNode root) {
        // Write your code here
        if (root == null) return null;
        if (root.val == 0 && root.left == null && root.right == null) return null;
        if (isValid(root.left)) {
            root.left = pruneTree(root.left);
        } else root.left = null;
        if (isValid(root.right)) {
            root.right = pruneTree(root.right);
        } else root.right = null;
        return root;
    }
    private boolean isValid(TreeNode node) {
        if (node == null) return false;
        if (node.val == 1 || isValid(node.left) ||isValid(node.right)) return true;
        else return false;
    }
}