【二叉树】二叉树剪枝

0x00 题目

给你二叉树的根结点 root
此外树的每个结点的值要么是 0，要么是 1
返回移除了所有不包含 1 的子树的原二叉树
节点 node 的子树为 node 本身
加上所有 node 的后代

---

0x01 思路

叶子节点值为 0 时，去掉
某个节点的如果要去掉
则左子树的值全为 0
右子树的值是全为 0
再加上节点本身的值也是 0
反过来讲，也就是
以某个节点为根的子树
只要存在值为 1 的节点
则这棵子树不用删除

---

0x02 解法

语言：Swift

树节点：TreeNode

public class TreeNode {
    public var val: Int
    public var left: TreeNode?
    public var right: TreeNode?
    public init() { self.val = 0; self.left = nil; self.right = nil; }
    public init(_ val: Int) { self.val = val; self.left = nil; self.right = nil; }
    public init(_ val: Int, _ left: TreeNode?, _ right: TreeNode?) {
        self.val = val
        self.left = left
        self.right = right
    }
}

解法：

func pruneTree(_ root: TreeNode?) -> TreeNode? {
    func trim(_ root: TreeNode?) -> Bool {
        guard let root = root else {
            return false
        }

        let left = trim(root.left)
        let right = trim(root.right)
        
        // 因为要从叶子节点开始剪，所以使用后序遍历方式
        // 左子树中全是 0
        if left == false {
            root.left = nil
        }
        // 右子树中全是 0
        if right == false {
            root.right = nil
        }
        
        // 左子树，右子树，节点，整个子树是否全为 0
        return left || right || (root.val == 1)
    }
    
    // 整棵树都为 0，则返回 nil
    let b = trim(root)
    return b ? root : nil
}

---

0x03 我的作品

欢迎体验我的作品之一：小五笔86版
五笔学习好帮手
App Store 搜索即可~

---