题目:
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
or1
.
思路:
我们采取先分后合的递归策略:首先prune根节点root的左右子树,并分别返回prune后的结果,然后判断root是不是应该被prune掉(应该被prune掉的条件是root的左右子树都返回了NULL,并且root的值为0)。如果是,则删除掉root结点,并返回NULL;否则说明以root为根节点的树不能被prune掉,但是它的子树有可能可以被prune掉,所以我们就给root接上它被prune掉的左右子树,然后返回root本身即可。
算法的时间复杂度为O(n),其中n为二叉树中的节点个数;如果不考虑递归调用所占用的栈空间,则算法的空间复杂度为O(1)。
代码:
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode(int x) : val(x), left(NULL), right(NULL) {}
* };
*/
class Solution {
public:
TreeNode* pruneTree(TreeNode* root) {
if (root == NULL) {
return NULL;
}
TreeNode *left = pruneTree(root->left);
TreeNode *right = pruneTree(root->right);
if (left == NULL && right == NULL && root->val == 0) {
delete root;
return NULL;
}
root->left = left;
root->right = right;
return root;
}
};