[LeetCode] 663. Equal Tree Partition_Medium tag: DFS, divide and conquer

Given the root of a binary tree, return true if you can partition the tree into two trees with equal sums of values after removing exactly one edge on the original tree.

 

Example 1:

Input: root = [5,10,10,null,null,2,3]
Output: true

Example 2:

Input: root = [1,2,10,null,null,2,20]
Output: false
Explanation: You cannot split the tree into two trees with equal sums after removing exactly one edge on the tree.

 

Constraints:

• The number of nodes in the tree is in the range [1, 104].
• -105 <= Node.val <= 105

 

Ideas: 实际就是求有没有subtree有total sum/2.0,

Note: 不能是root, 所以要将root 的sum 给去掉, below is an edge case

         0

     /      \

   -1       1

T: O(n)       S: O(n)

1. 建helper function, 和ans来得到从bottom 到up的subtree 的sum

2. total = ans.pop()

3. 最后将total/2.0 在ans里面找就好，因为T: O(n), 所以最后将ans 转换为set意义不大。

 

Code:

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def checkEqualTree(self, root: Optional[TreeNode]) -> bool:
        ans = []
        self.treeSum(root, ans)
        total = ans.pop()
        return total/2.0 in ans
    
    
    def treeSum(self, root, ans):
        if not root: return 0
        ans.append(root.val + self.treeSum(root.left, ans) + self.treeSum(root.right, ans))
        return ans[-1]