题目:
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; } };