[Swift]LeetCode814. 二叉树剪枝 | Binary Tree Pruning

★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★
➤微信公众号：山青咏芝（shanqingyongzhi）
➤博客园地址：山青咏芝（https://www.cnblogs.com/strengthen/）
➤GitHub地址：https://github.com/strengthen/LeetCode
➤原文地址： https://www.cnblogs.com/strengthen/p/10564091.html 
➤如果链接不是山青咏芝的博客园地址，则可能是爬取作者的文章。
➤原文已修改更新！强烈建议点击原文地址阅读！支持作者！支持原创！
★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★

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 or 1.

---

给定二叉树根结点 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 。

---

Runtime: 8 ms

Memory Usage: 19.3 MB

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     public var val: Int
 *     public var left: TreeNode?
 *     public var right: TreeNode?
 *     public init(_ val: Int) {
 *         self.val = val
 *         self.left = nil
 *         self.right = nil
 *     }
 * }
 */
class Solution {
    func pruneTree(_ root: TreeNode?) -> TreeNode? {
        if root == nil {return nil}
        root?.left = pruneTree(root?.left)
        root?.right = pruneTree(root?.right)
        return (root?.left == nil && root?.right == nil && root!.val == 0) ? nil : root
    }
}

---

12ms

class Solution {
    func pruneTree(_ root: TreeNode?) -> TreeNode? {
        _ = helper(root)
        return root
    }

    func helper(_ root: TreeNode?) -> Int {
        guard let root = root else {
            return 0 
        }
        let left = helper(root.left)
        let right = helper(root.right)
        if left == 0 {
            root.left = nil 
        }
        if right == 0 {
            root.right = nil 
        }
        return root.val + left + right
    }    
}

---

16ms

class Solution {
    private class item {
        var node: TreeNode
        var toVisit: Bool
        var parent: TreeNode
        var dir: Bool //left: false, right: true
        
        init(_ node: TreeNode, _ toVisit: Bool, _ parent: TreeNode, _ dir: Bool) {
            self.node = node
            self.toVisit = toVisit
            self.parent = parent
            self.dir = dir
        }
    }
    
    func pruneTree(_ root: TreeNode?) -> TreeNode? {
        guard let root = root else { return nil }
        let fake = TreeNode(-1); fake.left = root
        
        var stack = [item(root,true,fake,false)]
        
        while !stack.isEmpty {
            let top = stack.last!
            if top.toVisit {
                let tmp = top.node
                if (tmp.left != nil) {
                    stack.append(item(tmp.left!,true,tmp,false))
                }
                if (tmp.right != nil) {
                    stack.append(item(tmp.right!,true,tmp,true))
                }
                top.toVisit = false
            } else {
                let tmp = top.node
                if (tmp.left == nil && tmp.right == nil && tmp.val == 0) {
                    if top.dir {
                        top.parent.right = nil
                    } else {
                        top.parent.left = nil
                    }
                }
                stack.removeLast(1)
            }
        }
        
        return fake.left
    }
}