Leetcode 0663 Equal Tree Partition

Problem statement

https://leetcode.com/problems/equal-tree-partition/

Solution

Actually what we need to find is there exists subtree of our tree with sum of values equal to sum of all values, divided by 2. So, let us use function SumAll(node, v), which traverse our tree and also check if sum of current subtree is equal to v. On first run we just evaluate sum of all values S in tree. In second run we find value S/2 in our subtree.

Complexity

Time complexity is O(n), space complexity is O(h). (if we use bfs we have O(w) space complexity).

Code

class Solution:
    def checkEqualTree(self, root):
        def SumAll(node, v):
            if not node: return 0
            curr = SumAll(node.left, v) + node.val + SumAll(node.right, v)
            if curr == v and node != self.root: self.Found = True
            return curr
            
        self.Found = False
        self.root = root
        S = SumAll(root, float("inf"))
        T = SumAll(root, S/2)
        return self.Found