Problem Description

The Equal Tree Partition problem asks you to determine whether it is possible to split a given binary tree into two trees with equal sum by removing exactly one edge. You are given the root node of a binary tree, where each node contains an integer value. Your task is to check if there exists a way to remove one edge (i.e., cut the tree into two subtrees) such that the sum of the values of the nodes in both resulting trees is the same.

• The tree is non-empty and may contain positive, negative, or zero values.
• You may not reuse nodes—each belongs to exactly one subtree after the cut.
• There may be at most one valid way to partition the tree as described.
• The input is provided as a TreeNode object representing the root of the binary tree.
• The output is a boolean: true if such a partition exists, false otherwise.

Thought Process

To solve this problem, let's start by understanding what it means to partition the tree. If we remove an edge, we get two subtrees. For the partition to be "equal," the sum of all node values in each subtree must be the same.

The brute-force way would be to try every possible edge removal, calculate the sum of both resulting trees, and check if they are equal. However, this is inefficient because recalculating subtree sums repeatedly is time-consuming.

Instead, we notice that if the total sum of the tree is T, then after removing an edge, the sums of the two resulting trees must both be T/2. So, the problem reduces to checking if there exists a subtree (other than the whole tree itself) whose sum is exactly T/2.

This leads us to a more efficient approach: traverse the tree once to compute the total sum, then traverse again to check for a subtree with sum T/2.

Solution Approach

• Step 1: Compute Total Sum
  • Perform a post-order traversal of the tree to compute the sum of all nodes. Let's call this totalSum.
• Step 2: Check for Partition Possibility
  • If totalSum is odd, it's impossible to split into two equal integers, so return false.
  • Otherwise, the target sum for each partition is totalSum / 2.
• Step 3: Find Subtree with Target Sum
  • Traverse the tree again (or during the first traversal), and for each subtree, calculate its sum.
  • If any subtree (except the entire tree itself) has sum equal to totalSum / 2, return true.
  • Use a set or list to store all subtree sums for easy checking.
• Why This Works:
  • Removing an edge is equivalent to "cutting off" a subtree. If that subtree's sum is half the total, the remaining tree will also have half the total sum.

Example Walkthrough

Consider the following example tree:

        5
       / \
      10 10
         / \
        2   3
  

• Step 1: Compute the total sum:
  5 + 10 + 10 + 2 + 3 = 30
• Step 2: Target sum for each partition is 30 / 2 = 15
• Step 3: Traverse to find a subtree with sum 15:
  • The left subtree rooted at the left child (10) has sum 10.
  • The right subtree rooted at the right child (10) has sum 10 + 2 + 3 = 15.
• Step 4: Since we found a subtree (the right child) with sum 15, we can partition the tree by removing the edge between the root and this right child.
• Result: Return true.

Time and Space Complexity

• Brute-force Approach:
  • For every edge, compute the sum of both resulting trees. This is O(N^2) time, where N is the number of nodes.
  • Space is O(N) for recursion stack.
• Optimized Approach:
  • We traverse the tree once to compute the total sum (O(N)), and again to check for the partition (O(N)), so total time is O(N).
  • Space is O(N) for storing subtree sums and recursion stack.

Summary

The Equal Tree Partition problem is elegantly solved by reducing it to checking if any subtree sums to half the total tree sum. By leveraging post-order traversal and set storage, we avoid redundant computations and achieve linear time complexity. The key insight is recognizing that removing an edge is equivalent to isolating a subtree, and the sum check naturally follows. This approach is both efficient and easy to implement, making it a robust solution for binary tree partitioning challenges.

Code Implementation

# Definition for a binary tree node.
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Solution:
    def checkEqualTree(self, root: TreeNode) -> bool:
        subtree_sums = []
        
        def get_sum(node):
            if not node:
                return 0
            s = node.val + get_sum(node.left) + get_sum(node.right)
            subtree_sums.append(s)
            return s
        
        total = get_sum(root)
        # Remove the total sum (whole tree) from consideration
        subtree_sums.pop()
        if total % 2 != 0:
            return False
        return (total // 2) in subtree_sums
      

// Definition for a binary tree node.
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

class Solution {
public:
    bool checkEqualTree(TreeNode* root) {
        vector<int> subtreeSums;
        int total = getSum(root, subtreeSums);
        // Remove the sum of the whole tree
        subtreeSums.pop_back();
        if (total % 2 != 0) return false;
        for (int s : subtreeSums) {
            if (s == total / 2) return true;
        }
        return false;
    }
    
private:
    int getSum(TreeNode* node, vector<int>& sums) {
        if (!node) return 0;
        int s = node->val + getSum(node->left, sums) + getSum(node->right, sums);
        sums.push_back(s);
        return s;
    }
};
      

// Definition for a binary tree node.
public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { val = x; }
}

class Solution {
    private List<Integer> subtreeSums = new ArrayList<>();

    public boolean checkEqualTree(TreeNode root) {
        int total = getSum(root);
        // Remove the sum of the whole tree
        subtreeSums.remove(subtreeSums.size() - 1);
        if (total % 2 != 0) return false;
        return subtreeSums.contains(total / 2);
    }

    private int getSum(TreeNode node) {
        if (node == null) return 0;
        int s = node.val + getSum(node.left) + getSum(node.right);
        subtreeSums.add(s);
        return s;
    }
}
      

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = (val===undefined ? 0 : val)
 *     this.left = (left===undefined ? null : left)
 *     this.right = (right===undefined ? null : right)
 * }
 */
/**
 * @param {TreeNode} root
 * @return {boolean}
 */
var checkEqualTree = function(root) {
    const subtreeSums = [];
    function getSum(node) {
        if (!node) return 0;
        const s = node.val + getSum(node.left) + getSum(node.right);
        subtreeSums.push(s);
        return s;
    }
    const total = getSum(root);
    subtreeSums.pop(); // remove the total sum (whole tree)
    if (total % 2 !== 0) return false;
    return subtreeSums.includes(total / 2);
};