814_binary_tree_pruning

814. 二叉树剪枝 中等

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

代码参考：

package main

import "fmt"

func main() {
    root := &TreeNode{Val: 1}
    root.Right = &TreeNode{Val: 0}
    root.Right.Left = &TreeNode{Val: 0}
    root.Right.Right = &TreeNode{Val: 1}
    newRoot := pruneTree(root)
    fmt.Println(newRoot)
    fmt.Println(newRoot.Left)
    fmt.Println(newRoot.Right)
    fmt.Println(newRoot.Right.Left)
    fmt.Println(newRoot.Right.Right) // bingo
}

// 其实就是一个后序遍历，一步一步摘除 0  叶子
func pruneTree(root *TreeNode) *TreeNode {
    if root == nil {
        return nil
    }
    root.Left = pruneTree(root.Left)
    root.Right = pruneTree(root.Right)
    if root.Left == nil && root.Right == nil && root.Val == 0 {
        return nil
    }
    return root
}