【Lintcode】1003. Binary Tree Pruning

题目地址：

https://www.lintcode.com/problem/binary-tree-pruning/description

给定一棵二叉树，每个节点的值都是$0$或者$1$。要求将值全为$0$的子树都从二叉树中删掉，并返回删后的树根。

对于一个以root为根二叉树而言，将其值全为$0$的子树删掉等价于，先将其左右子树中值全为$0$的子树都删掉，再看root自己是否应该被删掉。很显然，如果root的左右子树都被删干净了，并且root自己值为$0$，说明这个以root为根的二叉树所有节点都是$0$，那么这个root自己也应该被删掉。

由这个思路，实现起来可以用递归。如果树本身是空树，就直接返回null；否则就先递归删除左右子树中值全为$0$的子树，并接到root上，再判断root的左右子树是否被删干净了，如果是，并且root本身也是$0$，则说明root本身也要删掉，则返回null；否则返回root即可。代码如下：

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;
        }
        // 递归删除左右子树，并接到root两边
        root.left = pruneTree(root.left);
        root.right = pruneTree(root.right);
        // 如果root本身也要删，则返回null
        if (root.left == null && root.right == null && root.val == 0) {
            return null;
        } else {
        	// 否则返回root
        	return root;
        }
    }
}

class TreeNode {
    int val;
    TreeNode left, right;
    TreeNode(int x) {
        val = x;
    }
}

时间复杂度$O(n)$，空间$O(h)$。