Equal Tree Partition

Problem Id: 663 Difficulty: Medium Tag: Tree


Intuition

Let's suggest that the edge we removed is e. The upper node of e is u and the lower node of e is l. Then the sub-tree with l as root node is one partition. And its sum should be 1/2 of the total sum of the original tree.

Solution


# 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:
        partition_sum = self._dfs_sum(root) / 2
        self.exist = False
        self._check_exist(root, partition_sum, isroot=True)
        return self.exist

    def _dfs_sum(self, node):
        if node is None:
            return 0
        return self._dfs_sum(node.left) + self._dfs_sum(node.right) + node.val

    def _check_exist(self, node, partition_sum, isroot=False):
        if node is None:
            return 0
        if self.exist:
            return 0
        current = node.val
        current += self._check_exist(node.left, partition_sum)
        current += self._check_exist(node.right, partition_sum)
        if current == partition_sum and not isroot:
            self.exist = True
        return current