663. Equal Tree Partition

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.

Example 1:

Input:     
    5
   / \
  10 10
    /  \
   2   3

Output: True
Explanation: 
    5
   / 
  10
      
Sum: 15

   10
  /  \
 2    3

Sum: 15

Example 2:

Input:     
    1
   / \
  2  10
    /  \
   2   20

Output: False
Explanation: You can't split the tree into two trees with equal sum after removing exactly one edge on the tree.

Note:

  1. The range of tree node value is in the range of [-100000, 100000].
  2. 1 <= n <= 10000

Rust Solution

struct Solution;
use rustgym_util::*;

trait Postorder {
    fn sum_size(&self) -> (i32, usize);
    fn check(&self, found: &mut bool, sum: i32, size: usize) -> (i32, usize);
}

impl Postorder for TreeLink {
    fn sum_size(&self) -> (i32, usize) {
        if let Some(node) = self {
            let node = node.borrow();
            let val = node.val;
            let left = node.left.sum_size();
            let right = node.right.sum_size();
            (val + left.0 + right.0, 1 + left.1 + right.1)
        } else {
            (0, 0)
        }
    }
    fn check(&self, found: &mut bool, sum: i32, size: usize) -> (i32, usize) {
        if let Some(node) = self {
            let node = node.borrow();
            let val = node.val;
            let left = node.left.check(found, sum, size);
            let right = node.right.check(found, sum, size);
            let res = (val + left.0 + right.0, 1 + left.1 + right.1);
            if res.0 * 2 == sum && res.1 != size {
                *found = true;
            }
            res
        } else {
            (0, 0)
        }
    }
}

impl Solution {
    fn check_equal_tree(root: TreeLink) -> bool {
        let (sum, size) = root.sum_size();
        let mut res = false;
        root.check(&mut res, sum, size);
        res
    }
}

#[test]
fn test() {
    let root = tree!(5, tree!(10), tree!(10, tree!(2), tree!(3)));
    let res = true;
    assert_eq!(Solution::check_equal_tree(root), res);
}

Having problems with this solution? Click here to submit an issue on github.