938. Range Sum of BST - cocoder39/coco39_LC GitHub Wiki

938. Range Sum of BST

Solution 1: dfs

do root check in dfs instead of before dfs to simplify the code structure. Otherwise there will be none node check everywhere

class Solution:
    def rangeSumBST(self, root: Optional[TreeNode], low: int, high: int) -> int:
        res = [0]
        
        def dfs(root) -> int:
            if root:
                if low <= root.val <= high:
                    res[0] += root.val

                if root.val < low:
                    dfs(root.right)
                elif root.val > high:
                    dfs(root.left)
                else: # low <= root.val <= high
                    dfs(root.left)
                    dfs(root.right)

        if root:
            dfs(root)
        return res[0] 

improvement: * consolidate branches in dfs

class Solution:
    def rangeSumBST(self, root: Optional[TreeNode], low: int, high: int) -> int:
        def dfs(root) -> int:
            res = 0
            if root:
                if low <= root.val <= high:
                    res += root.val

                if root.val <= high:
                    res += dfs(root.right)
                if root.val >= low:
                    res += dfs(root.left)
            return res

        return dfs(root)

solution 2: iterative solution using stack

class Solution:
    def rangeSumBST(self, root: Optional[TreeNode], low: int, high: int) -> int:
        res = 0
        stack = [root]
        while stack:
            node = stack.pop()
            if node:
                if low <= node.val <= high:
                    res += node.val
                if node.val > low: # value smaller than node.val should be considered
                    stack.append(node.left) 
                if node.val < high: # value bigger than node.val should be considered
                    stack.append(node.right)
        return res