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