Maximum Product of Splitted Binary Tree LeetCode solution in Python

1
class Solution:
2
    def maxProduct(self, root: Optional[TreeNode]) -> int:
3
        def findTotalSum(node, totalSum):
4
            if node is None:
5
                return totalSum
6
            totalSum = findTotalSum(node.left, totalSum)
7
            totalSum += node.val
8
            totalSum = findTotalSum(node.right, totalSum)
9
            return totalSum
10

11
        def dfs(node, maxProd, totalSum):
12
            if node is None:
13
                return maxProd, 0
14
            if not node.left and not node.right:
15
                return maxProd, node.val
16
            maxProd, lSum = dfs(node.left, maxProd, totalSum)
17
            maxProd, rSum = dfs(node.right, maxProd, totalSum)
18
            subTreeSum = lSum + rSum + node.val
19
            maxProd = max(
20
                maxProd,
21
                (totalSum - lSum) * lSum,
22
                (totalSum - rSum) * rSum,
23
                (totalSum - subTreeSum) * subTreeSum,
24
            )
25
            return maxProd, subTreeSum
26

27
        totalSum = findTotalSum(root, 0)
28
        product, _ = dfs(root, 1, totalSum)
29
        return product % (pow(10, 9) + 7)