Leetcode 663. Equal Tree Partition

Problem Explanation

Given a binary tree we have to figure out whether we can partition the tree into two sub-trees that have the same sum. Partitioning of the tree is done by removing exactly one edge.

Let's take an example to understand this:

Consider the tree:

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

The sum of the values in the whole tree is 30 (5+10+10+2+3). If we remove the edge between the root and the right child, we get two separate trees, one having the root as 5 and the other having the root as 10. The sum of values in both these newly formed trees is 15 (5+10) and 15 (10+2+3) respectively. These sums are the same, thus it is possible to partition the original tree into two trees having the same sum of values.

Solution Approach

This problem can be solved with the help of Depth First Search (DFS). During DFS, we calculate the sum of all nodes. While backtracking, we store the sum of subtree rooted at the current node, in the hashset. If we find a sum that is half of the total sum, we return true.

Below is the solution to the problem in different programming languages.

Python Solution

1
2python
3class Solution:
4    def checkEqualTree(self, root):
5        seen = []
6        
7        def sum_(node):
8            if not node:
9                return 0
10            seen.append(sum_(node.left) + sum_(node.right) + node.val)
11            return seen[-1]
12        
13        total = sum_(root)
14        seen.pop()
15        return total / 2.0 in seen

Java Solution

1
2java
3class Solution {
4    public boolean checkEqualTree(TreeNode root) {
5        Set<Integer> set = new HashSet();
6        int sum = sum(root, set);
7        return sum % 2 == 0 && set.contains(sum / 2);
8    }
9
10    public int sum(TreeNode root, Set<Integer> set) {
11        if (root == null) {
12            return 0;
13        }
14        int s = root.val + sum(root.left, set) + sum(root.right, set);
15        if (root != root) {
16            set.add(s);
17        }
18        return s;
19    }
20}

JavaScript Solution

1
2javascript
3var checkEqualTree = function(root) {
4    let map = {}, isPossible = false;
5    const total = getNodeVal(root);
6    if(total == 0) return map[0] > 1;
7    return total%2 == 0 && !!map[total/2];
8
9    
10    function getNodeVal(node) {
11        if(!node) return 0;
12        let sum = node.val + getNodeVal(node.left) + getNodeVal(node.right);
13        map[sum] = (map[sum] || 0) + 1;
14        return sum;
15    }
16};

C++ Solution

1
2c++
3class Solution {
4public:
5    bool checkEqualTree(TreeNode* root) {
6        int total = sum(root);
7        if (total % 2 == 0)
8            return dfs(root, total / 2, 0);
9        return false;
10    }
11private:
12    int sum(TreeNode* root) {
13        if (root == NULL)
14            return 0;
15        return root -> val + sum(root -> left) + sum(root -> right);
16    }
17    
18    bool dfs(TreeNode* root, int target, int pre) {
19        if (root == NULL)
20            return false;
21        pre += root -> val;
22        if (pre == target)
23            return true;
24        return dfs(root -> left, target, pre) || dfs(root -> right, target, pre);
25    }
26};

C# Solution

1
2csharp
3public class Solution 
4{
5    public bool CheckEqualTree(TreeNode root) 
6    {
7        HashSet<int> sums = new HashSet<int>();
8        int sum = GetSum(root, sums);
9        if(sum%2 == 0) 
10        {
11            return sums.Contains(sum/2);
12        }
13        else 
14            return false;
15    }
16    
17    private int GetSum(TreeNode node, HashSet<int>sums) 
18    {
19        if(node == null) return 0;
20        int sum = node.val + GetSum(node.left, sums) + GetSum(node.right, sums);
21        sums.Add(sum);
22        return sum;
23    }
24}

Golang Solution

A solution in Golang can be approached by creating a helper method ```

sumTree```

which finds the sum of all nodes in a tree and returns it as an int. This is used by the main function ```

checkEqualTree```

.

1
2go
3type TreeNode struct {
4     Val int
5     Left *TreeNode
6     Right *TreeNode
7}
8
9func checkEqualTree(root *TreeNode) bool {
10    sums := make(map[int]bool)
11    total := sumTree(root, sums)
12    _, ok := sums[total/2]
13    return ok && total%2 == 0
14}
15
16func sumTree(node *TreeNode, sums map[int]bool) int {
17    if node == nil {
18        return 0
19    }
20    sum := node.Val + sumTree(node.Left, sums) + sumTree(node.Right, sums)
21    sums[sum] = true
22    return sum
23}

Rust Solution

In Rust, we can accomplish this by creating a helper function to sum up the node values while keeping track of the sums so far. The ```

check_equal_tree```

function returns true if the total sum is divisible by 2 and the sum that is half of the total sum is in the list of sums.

1
2rust
3# Definition for a binary tree node.
4# struct TreeNode {
5#     val: i32,
6#     left: Option<Rc<RefCell<TreeNode>>>,
7#     right: Option<Rc<RefCell<TreeNode>>>,
8# }
9
10use std::rc::Rc;
11use std::cell::RefCell;
12
13impl Solution {
14    fn dfs(node: &Option<Rc<RefCell<TreeNode>>>, sums: &mut Vec<i32>) -> i32{
15        if let Some(n) = node{
16            let left = Self::dfs(&n.borrow().left, sums);
17            let right = Self::dfs(&n.borrow().right, sums);
18            let sum = left + right + n.borrow().val;
19            sums.push(sum);
20            return sum;
21        }
22        0
23    }
24
25    pub fn check_equal_tree(root: Option<Rc<RefCell<TreeNode>>>) -> bool {
26        let mut sums = vec![];
27        let root_val = Self::dfs(&root, &mut sums);
28        sums.pop();
29
30        root_val % 2 == 0 && sums.contains(&(root_val / 2))
31    }
32}

This code uses Rust's standard Option and Rc<RefCell<T>> types to safely handle possible null values and shared mutable state, respectively.