（Java）leetcode-814 Binary Tree Pruning（二叉树剪枝）

题目描述

给定二叉树根结点 root ，此外树的每个结点的值要么是 0，要么是 1。

返回移除了所有不包含 1 的子树的原二叉树。

( 节点 X 的子树为 X 本身，以及所有 X 的后代。)

示例1:
输入: [1,null,0,0,1]
输出: [1,null,0,null,1]

解释:
只有红色节点满足条件“所有不包含 1 的子树”。
右图为返回的答案。

示例2:
输入: [1,0,1,0,0,0,1]
输出: [1,null,1,null,1]

示例3:
输入: [1,1,0,1,1,0,1,0]
输出: [1,1,0,1,1,null,1]

说明:

给定的二叉树最多有 100 个节点。
每个节点的值只会为 0 或 1 。

思路

剪掉一棵子树的条件是“子树不包含1”，判断是否该剪枝，只需要判断这课树是否含“1”。
那么可以将问题分解为：
一棵树包含“1” = 根节点含“1”或左子树含“1”或右子树“含1”
依照上面的关系就能通过 后序遍历 写出 自下而上 的解决方案。先判断左右子树是否含1，最后通过根节点的值判断整棵树是否含1。剪枝动作可以在判断完左右子树后就执行。

时间复杂度：O(n) 每个节点只访问一遍
空间复杂度：O(n) 递归时使用栈的开销

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

    private boolean hasOne(TreeNode root) {
    	if (root == null) return false;
    	// 递归判断左右子树是否含1
    	boolean left = hasOne(root.left);
    	boolean right = hasOne(root.right); 

    	// 执行剪枝
    	if (!left) root.left = null;
    	if (!right) root.right = null;

    	return root.val == 1 || left || right;
    }
}

执行用时：0 ms, 在所有 Java 提交中击败了100.00%的用户
内存消耗：37.6 MB, 在所有 Java 提交中击败了8.33%的用户

如果追求极度的精简，代码可以写成这样，逻辑是一样的：

class Solution {
    public TreeNode pruneTree(TreeNode root) {
        if (root == null) return null;
        root.left = pruneTree(root.left);
        root.right = pruneTree(root.right);
        return (root.val == 1 || root.left != null || root.right != null) ? root : null;
    }
}