（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;
    }
}