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

★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★
➤微信公衆號：山青詠芝（shanqingyongzhi）
➤博客園地址：山青詠芝（https://www.cnblogs.com/strengthen/）
➤GitHub地址：https://github.com/strengthen/LeetCode
➤原文地址： http://www.javashuo.com/article/p-zlwrsmgd-me.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

 1 /**
 2  * Definition for a binary tree node.
 3  * public class TreeNode {
 4  *     public var val: Int
 5  *     public var left: TreeNode?
 6  *     public var right: TreeNode?
 7  *     public init(_ val: Int) {
 8  *         self.val = val
 9  *         self.left = nil
10  *         self.right = nil
11  *     }
12  * }
13  */
14 class Solution {
15     func pruneTree(_ root: TreeNode?) -> TreeNode? {
16         if root == nil {return nil}
17         root?.left = pruneTree(root?.left)
18         root?.right = pruneTree(root?.right)
19         return (root?.left == nil && root?.right == nil && root!.val == 0) ? nil : root
20     }
21 }

---

12ms

 1 class Solution {
 2     func pruneTree(_ root: TreeNode?) -> TreeNode? {
 3         _ = helper(root)
 4         return root
 5     }
 6 
 7     func helper(_ root: TreeNode?) -> Int {
 8         guard let root = root else {
 9             return 0 
10         }
11         let left = helper(root.left)
12         let right = helper(root.right)
13         if left == 0 {
14             root.left = nil 
15         }
16         if right == 0 {
17             root.right = nil 
18         }
19         return root.val + left + right
20     }    
21 }

---

16ms

 1 class Solution {
 2     private class item {
 3         var node: TreeNode
 4         var toVisit: Bool
 5         var parent: TreeNode
 6         var dir: Bool //left: false, right: true
 7         
 8         init(_ node: TreeNode, _ toVisit: Bool, _ parent: TreeNode, _ dir: Bool) {
 9             self.node = node
10             self.toVisit = toVisit
11             self.parent = parent
12             self.dir = dir
13         }
14     }
15     
16     func pruneTree(_ root: TreeNode?) -> TreeNode? {
17         guard let root = root else { return nil }
18         let fake = TreeNode(-1); fake.left = root
19         
20         var stack = [item(root,true,fake,false)]
21         
22         while !stack.isEmpty {
23             let top = stack.last!
24             if top.toVisit {
25                 let tmp = top.node
26                 if (tmp.left != nil) {
27                     stack.append(item(tmp.left!,true,tmp,false))
28                 }
29                 if (tmp.right != nil) {
30                     stack.append(item(tmp.right!,true,tmp,true))
31                 }
32                 top.toVisit = false
33             } else {
34                 let tmp = top.node
35                 if (tmp.left == nil && tmp.right == nil && tmp.val == 0) {
36                     if top.dir {
37                         top.parent.right = nil
38                     } else {
39                         top.parent.left = nil
40                     }
41                 }
42                 stack.removeLast(1)
43             }
44         }
45         
46         return fake.left
47     }
48 }