LeetCode ---- 二叉树剪枝

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

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
/**
 * -------------------------------------------------<br>
 * Copyright (c) 2015-2017 AIA . All Rights Reserved.  <br>
 * -------------------------------------------------<br>
 * Project Name   : <code>LeetCode</code><br>
 *
 * @author: Tal.Yuan
 * @version: 2.0
 * Description:
 */
package com.tal.leetcode;

import org.junit.Assert;
import org.junit.Test;

import java.util.Stack;


public class PruneTree {

    @Test
    public void pruneTreeTest() {
        TreeNode test = new TreeNode(1);
        test.right = new TreeNode(0);
        test.right.left = new TreeNode(0);
        test.right.right = new TreeNode(1);

        TreeNode actual = pruneTree(test);

        Assert.assertEquals(1, actual.val);
        Assert.assertNull(actual.left);
        Assert.assertEquals(0, actual.right.val);
        Assert.assertNull(actual.right.left);
        Assert.assertEquals(1, actual.right.right.val);
    }

    public TreeNode pruneTree(TreeNode root) {
        TreeNode temp = root;
        Stack<TreeNode> stack = new Stack<>();
        while (root != null || stack.size() > 0) {
            while (root != null) {
                if (root.left != null && checkInvalidNode(root.left)) {
                    root.left = null;
                }

                if (root.right != null) {
                    stack.push(root);
                }

                root = root.left;
            }

            if (!stack.isEmpty()) {
                TreeNode treeNode = stack.pop();
                if (null != treeNode && checkInvalidNode(treeNode.right)) {
                    treeNode.right = null;
                }
                root = treeNode.right;
            }
        }

        return temp;
    }

    public boolean checkInvalidNode(TreeNode root) {
        Stack<TreeNode> stack = new Stack<>();
        while (root != null || !stack.isEmpty()) {
            while (root != null) {
                if (root.val == 1) {
                    return false;
                }
                stack.push(root);
                root = root.left;
            }

            if (!stack.isEmpty()) {
                TreeNode treeNode = stack.pop();
                if (null != treeNode && null != treeNode.right && treeNode.right.val == 1) {
                    return false;
                }
                root = treeNode.right;
            }
        }
        return true;
    }

}