Range sum in BST

range_sum_bst.py

class Node(object):

    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None
        # Store the range of values
        self.range = None
        # Store the sum of the tree planted at this node
        self.sum = None


class Solution(object):

    def pre_compute_tree(self, node):

        if not node:
            return 0, ()

        left = self.pre_compute_tree(node.left)
        right = self.pre_compute_tree(node.right)

        left_sum = left[0]
        if len(left[1]) > 0:
            left_min = left[1][0]
        else:
            left_min = None

        right_sum = right[0]
        if len(right[1]) > 0:
            right_max = right[1][1]
        else:
            right_max = None

        range_start = left_min or node.value
        range_end = right_max or node.value
        node_range = (range_start, range_end)
        node_sum = left_sum + right_sum + node.value

        node.range = node_range
        node.sum = node_sum

        return node.sum, node.range

    def solve(self, node, range_start, range_end):
        return self.helper(node, range_start, range_end)

    def helper(self, node, range_start, range_end):

        if not node:
            return 0

        node_range = node.range

        if node_range[0] >= range_start and node_range[1] <= range_end:
            return node.sum

        if node.right and node_range[1] <= range_end:
            right_sum = node.right.sum
        else:
            right_sum = self.helper(node.right, range_start, range_end)

        if node.left and node_range[0] >= range_start:
            left_sum = node.left.sum
        else:
            left_sum = self.helper(node.left, range_start, range_end)

        current_sum = left_sum + right_sum

        if range_start <= node.value <= range_end:
            current_sum += node.value

        return current_sum


if __name__ == "__main__":
    root = Node(15)
    root.left = Node(10)
    root.right = Node(20)
    root.left.left = Node(8)
    root.left.right = Node(12)
    root.right.left = Node(16)

    Solution().pre_compute_tree(root)
    ans = Solution().solve(root, 10, 16)
    print(ans)