Prefix Sum

Build a prefix sum array so any range sum can be answered in O(1). prefix[i] stores the sum of elements from index 0 to i-1, with prefix[0] = 0 as the base case.

from itertools import accumulate

def prefix_sum(nums: list[int]) -> list[int]:
    return [0] + list(accumulate(nums))

def range_sum(prefix: list[int], left: int, right: int) -> int:
    return prefix[right + 1] - prefix[left]

if __name__ == '__main__':
    nums = [1, 2, 3, 4, 5]
    prefix = prefix_sum(nums)
    print(prefix)  # [0, 1, 3, 6, 10, 15]

    print(range_sum(prefix, 1, 3))  # 9 (2+3+4)
    print(range_sum(prefix, 0, 4))  # 15 (1+2+3+4+5)
    print(range_sum(prefix, 2, 2))  # 3