[LintCode] Problem 864 - Equal Tree Partition

Posted on | In ,

Given a binary tree with n nodes, your task is to check if it’s possible to partition the tree to two trees which have the equal sum of values after removing exactly one edge on the original tree.

Note

  1. The range of tree node value is in the range of [-100000, 100000].
  2. 1 <= n <= 10000
  3. You can assume that the tree is not null

Example

No.1

Input: {5,10,10,#,#,2,3}

Output: true

Explanation:
origin:

1
2
3
4
5
  5
 / \
10 10
  /  \
 2    3

two subtrees:

1
2
3
  5       10
 /       /  \
10      2    3

No.2

Input: {1,2,10,#,#,2,20}

Output: false

Explanation:
origin:

1
2
3
4
5
  1
 / \
2  10
  /  \
 2    20

Code

1
2
3
4
5
6
7
8
public class TreeNode {
    public int val;
    public TreeNode left, right;
    public TreeNode(int val) {
        this.val = val;
        this.left = this.right = null;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
private Set<Integer> set = new HashSet<>();
private int countZero = 0;

public boolean checkEqualTree(TreeNode root) {
    int sum = helper(root);

    if (countZero != 0)
        return countZero % 2 == 0;
    
    return set.contains(sum >> 1);
}

private int helper(TreeNode root) {
    if (root == null)
        return 0;

    int left = helper(root.left);
    int right = helper(root.right);
    int sum = left + right + root.val;

    if (sum == 0)
        countZero++;

    set.add(sum);
    return sum;
}