Maximum Sum BST in Binary Tree LeetCode solution in Python

1
class Solution:
2
    def maxSumBST(self, root: TreeNode) -> int:
3
        res = 0
4

5
        def traverse(root):
6
            """return status_of_bst, size_of_bst, left_bound, right_bound"""
7
            nonlocal res
8
            if not root:
9
                return 1, 0, None, None  # this subtree is empty
10

11
            ls, l, ll, lr = traverse(root.left)
12
            rs, r, rl, rr = traverse(root.right)
13

14
            if ((ls == 2 and lr < root.val) or ls == 1) and (
15
                (rs == 2 and rl > root.val) or rs == 1
16
            ):
17
                # this subtree is a BST
18
                size = root.val + l + r
19
                res = max(res, size)
20
                return (
21
                    2,
22
                    size,
23
                    (ll if ll is not None else root.val),
24
                    (rr if rr is not None else root.val),
25
                )
26
            return 0, None, None, None  # this subtree is not a BST
27

28
        traverse(root)
29
        return res