Task

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:

1
2
3
4
5
6

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.

leetcode814_1

Example 2:

1
2

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

leetcode814_2

Example 3:

1
2

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

leetcode814_3

Note:

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

Solution

  • 只需递归遍历二叉树，从叶子节点开始删除，若一个节点值为0且无左右子树，则删除它。从叶子向根逐个删除节点，即可实现删除全为0的子树
  • 时间复杂度：每个节点遍历一次，O(n)
  • 空间复杂度：递归取决于树的高度，O(h)
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

// Runtime: 4 ms, faster than 64.22% of C++ online submissions for Binary Tree Pruning.
// Memory Usage: 8.6 MB, less than 6.49% of C++ online submissions for Binary Tree Pruning.
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        //走到叶子节点，递归终止
        if(!root) return nullptr;
        //遍历
        root->left=pruneTree(root->left);
        root->right=pruneTree(root->right);
        //值为1或存在左/右子树，则不删除
        if(root->val==1 || root->left || root->right)
            return root;
        //值为0且不存在左右子树，则删除
        delete root;
        return nullptr;
    }
};