| 
 | 1 | +from __future__ import annotations  | 
 | 2 | + | 
 | 3 | +import sys  | 
 | 4 | +from dataclasses import dataclass  | 
 | 5 | + | 
 | 6 | +INT_MIN = -sys.maxsize + 1  | 
 | 7 | +INT_MAX = sys.maxsize - 1  | 
 | 8 | + | 
 | 9 | + | 
 | 10 | +@dataclass  | 
 | 11 | +class TreeNode:  | 
 | 12 | +    val: int = 0  | 
 | 13 | +    left: TreeNode | None = None  | 
 | 14 | +    right: TreeNode | None = None  | 
 | 15 | + | 
 | 16 | + | 
 | 17 | +def max_sum_bst(root: TreeNode | None) -> int:  | 
 | 18 | +    """  | 
 | 19 | +    The solution traverses a binary tree to find the maximum sum of  | 
 | 20 | +    keys in any subtree that is a Binary Search Tree (BST). It uses  | 
 | 21 | +    recursion to validate BST properties and calculates sums, returning  | 
 | 22 | +    the highest sum found among all valid BST subtrees.  | 
 | 23 | +
  | 
 | 24 | +    >>> t1 = TreeNode(4)  | 
 | 25 | +    >>> t1.left = TreeNode(3)  | 
 | 26 | +    >>> t1.left.left = TreeNode(1)  | 
 | 27 | +    >>> t1.left.right = TreeNode(2)  | 
 | 28 | +    >>> print(max_sum_bst(t1))  | 
 | 29 | +    2  | 
 | 30 | +    >>> t2 = TreeNode(-4)  | 
 | 31 | +    >>> t2.left = TreeNode(-2)  | 
 | 32 | +    >>> t2.right = TreeNode(-5)  | 
 | 33 | +    >>> print(max_sum_bst(t2))  | 
 | 34 | +    0  | 
 | 35 | +    >>> t3 = TreeNode(1)  | 
 | 36 | +    >>> t3.left = TreeNode(4)  | 
 | 37 | +    >>> t3.left.left = TreeNode(2)  | 
 | 38 | +    >>> t3.left.right = TreeNode(4)  | 
 | 39 | +    >>> t3.right = TreeNode(3)  | 
 | 40 | +    >>> t3.right.left = TreeNode(2)  | 
 | 41 | +    >>> t3.right.right = TreeNode(5)  | 
 | 42 | +    >>> t3.right.right.left = TreeNode(4)  | 
 | 43 | +    >>> t3.right.right.right = TreeNode(6)  | 
 | 44 | +    >>> print(max_sum_bst(t3))  | 
 | 45 | +    20  | 
 | 46 | +    """  | 
 | 47 | +    ans: int = 0  | 
 | 48 | + | 
 | 49 | +    def solver(node: TreeNode | None) -> tuple[bool, int, int, int]:  | 
 | 50 | +        """  | 
 | 51 | +        Returns the maximum sum by making recursive calls  | 
 | 52 | +        >>> t1 = TreeNode(1)  | 
 | 53 | +        >>> print(solver(t1))  | 
 | 54 | +        1  | 
 | 55 | +        """  | 
 | 56 | +        nonlocal ans  | 
 | 57 | + | 
 | 58 | +        if not node:  | 
 | 59 | +            return True, INT_MAX, INT_MIN, 0  # Valid BST, min, max, sum  | 
 | 60 | + | 
 | 61 | +        is_left_valid, min_left, max_left, sum_left = solver(node.left)  | 
 | 62 | +        is_right_valid, min_right, max_right, sum_right = solver(node.right)  | 
 | 63 | + | 
 | 64 | +        if is_left_valid and is_right_valid and max_left < node.val < min_right:  | 
 | 65 | +            total_sum = sum_left + sum_right + node.val  | 
 | 66 | +            ans = max(ans, total_sum)  | 
 | 67 | +            return True, min(min_left, node.val), max(max_right, node.val), total_sum  | 
 | 68 | + | 
 | 69 | +        return False, -1, -1, -1  # Not a valid BST  | 
 | 70 | + | 
 | 71 | +    solver(root)  | 
 | 72 | +    return ans  | 
 | 73 | + | 
 | 74 | + | 
 | 75 | +if __name__ == "__main__":  | 
 | 76 | +    import doctest  | 
 | 77 | + | 
 | 78 | +    doctest.testmod()  | 
0 commit comments