Leetcode 1373. Maximum Sum BST in Binary Tree [Python]

初期版本TLE在第55个TC。判断BST的函数似乎得当作模版记下来。

# 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 maxSumBST(self, root: TreeNode) -> int:
        if not root:return 0
        curres = 0
        if self.isBST(root):
            curres = self.sumBST(root)
        return max(curres, max(self.maxSumBST(root.left), self.maxSumBST(root.right)))            
    def isBST(self, root):
        return self.isBSThelper(root, float('-inf'), float('inf'))
    def isBSThelper(self, root, curmin, curmax):
        if not root:return True
        if root.val <= curmin or root.val >= curmax:return False
        return self.isBSThelper(root.left, curmin, root.val) and self.isBSThelper(root.right, root.val, curmax)
    def sumBST(self, root):
        if not root:return 0
        return root.val + self.sumBST(root.left) + self.sumBST(root.right)      

改进一下

class Solution:
    def maxSumBST(self, root: TreeNode) -> int:
        self.res = float('-inf')
        def dfs(root):
            if not root: return float('inf'), float('-inf'), 0 #min,max,sum
            leftmin, leftmax, leftsum = dfs(root.left)
            rightmin, rightmax, rightsum = dfs(root.right)
            if leftmax < root.val < rightmin: # valid BST
                self.res = max(self.res, root.val + leftsum + rightsum)
                return min(root.val, leftmin), max(root.val , rightmax), root.val + leftsum + rightsum
            else:
                return float('-inf'), float('inf'), 0
        
        dfs(root)
        return max(self.res, 0)