二叉树剪枝(Java)

给定二叉树根结点?root?，此外树的每个结点的值要么是 0，要么是 1。
返回移除了所有不包含 1 的子树的原二叉树。
(节点 X 的子树为 X 本身，以及所有 X 的后代。)

示例1:
输入: [1,null,0,0,1]
输出: [1,null,0,null,1]
       1
         \
           0
         /    \
       0      1
输出后树为：
       1
         \
           0
             \
              1

示例2:
输入: [1,0,1,0,0,0,1]
输出: [1,null,1,null,1]
                1
             /     \
          0         1
        /    \      /    \
      0      0 0       1
输出后树为：
               1
                 \
                  1
                    \
                     1

示例3:
输入: [1,1,0,1,1,0,1,0]
输出: [1,1,0,1,1,null,1]
                1
              /   \
            1       0
          /  \     /    \
        1    1  0     1
       /
      0
输出后树为：
                1
              /   \
            1       0
          /  \        \
       1     1       1
         
        

说明:
给定的二叉树最多有 100 个节点。
每个节点的值只会为 0 或 1 。

package com.loo;

public class PruneTree {

    public static void main(String[] args) {
        TreeNode root = new TreeNode(1);
        TreeNode right1 = new TreeNode(0);
        TreeNode left11 = new TreeNode(0);
        TreeNode right12 = new TreeNode(1);
        root.right = right1;
        right1.left = left11;
        right1.right = right12;
        TreeNode pruneTree = getPruneTree(root);
        printTreeNode(pruneTree);
        System.out.println("++++++++++++++++++++");
        TreeNode rootB = new TreeNode(1);
        TreeNode leftB1 = new TreeNode(0);
        TreeNode rightB1 = new TreeNode(1);
        TreeNode leftB2 = new TreeNode(0);
        TreeNode leftB3 = new TreeNode(0);
        TreeNode rightB2 = new TreeNode(0);
        TreeNode rightB3 = new TreeNode(1);
        rootB.left = leftB1;
        rootB.right = rightB1;
        leftB1.left = leftB2;
        leftB1.right = leftB3;
        rightB1.left = rightB2;
        rightB1.right = rightB3;
        TreeNode pruneTreeB = getPruneTree2(rootB);
        printTreeNode(pruneTreeB);
    }
    
    public static TreeNode getPruneTree(TreeNode root) {
        return containsOne(root) ? root : null;
    }
    
    public static boolean containsOne(TreeNode root) {
        if (root == null) {
            return false;
        }
        boolean left = containsOne(root.left);
        boolean right = containsOne(root.right);
        if (!left) {
            root.left = null;
        }
        if (!right) {
            root.right = null;
        }
        return root.value == 1 || left || right;
    }
    
    // 从底向上后续遍历，如果当前元素为0且左右子树为null，则删除节点
    public static TreeNode getPruneTree2(TreeNode root) {
        if (root == null) {
            return null;
        }
        root.left = getPruneTree2(root.left);
        root.right = getPruneTree2(root.right);
        if (root.value == 0 && root.left == null && root.right == null) {
            return null;
        }
        return root;
    }
    
    public static void printTreeNode(TreeNode root) {
        if (root!=null) {
            System.out.println(root.value);
            printTreeNode(root.left);
            printTreeNode(root.right);
        }
    }
    
    static class TreeNode {
        int value;
        TreeNode left;
        TreeNode right;
        TreeNode(int v) {
            value = v;
        }
    }

}