算法 22周 Binary Tree Pruning

238 阅读1分钟

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.

暴力解法:

class Solution {
    public TreeNode pruneTree(TreeNode root) {
        if (root.left != null) {
            pruneTree(root.left);
            if(root.left.val==-1){
                root.left=null;
            }
        }

        if (root.right != null) {
            pruneTree(root.right);
            if(root.right.val==-1){
                root.right=null;
            }
        }

        if(root.left == null&&root.right == null&&root.val==0){
            root.val = -1;
        }
        
        return root;
    }
}

Runtime: 0 ms, faster than 100.00% of Java online submissions for Binary Tree Pruning. Memory Usage: 34.4 MB, less than 100.00% of Java online submissions for Binary Tree Pruning.

官方解法

class Solution {
    public TreeNode pruneTree(TreeNode root) {
        return containsOne(root) ? root : null;
    }

    public boolean containsOne(TreeNode node) {
        if (node == null) return false;
        boolean a1 = containsOne(node.left);
        boolean a2 = containsOne(node.right);
        if (!a1) node.left = null;
        if (!a2) node.right = null;
        return node.val == 1 || a1 || a2;
    }
}