Description

Prefix sum is a method to solve a particular type of problem where you need to read sum of a sub-array multiple times. A prefix sum is a separate array generally equal in length to actual array. Each index element of the prefix sum is sum of all element up to the index element of actual array.

Note that, in above diagram prefix-sum[index] = sum of elements of original array from 0 to index

Usage

Prefix sum can solve certain types of problems where sum of sub-array is required multiple times. For example:

1. Given a binary array, return the maximum length of a sub-array with an equal number of 0 and 1. Check out my blog
2. Calculate length of largest sub-array with sum equal to a certain value or zero
3. Calculate range sum for multiple (start, end) queries

Approach

Let’s try to come up with an optimized approach to calculate prefix-sum step by step.

Naive

A naive approach is to use 2 loops. Outer loop keeps track of index up to which we want to sum elements of original array. Inner loop goes from 0 to index to actually calculate sum.

def calculate_pre_sum_naive(nums: list[int]):
    prefix_sum = []
    for index in range(len(nums)):
        sum_index = 0
        for j in range(index + 1):
            sum_index += nums[j]
        prefix_sum.append(sum_index)
    
    return prefix_sum

Complexity

$\textbf{Time Complexity} = O(N^2)$

$\textbf{Space Complexity} = O(N)$

Optimized

We can optimize the above approach by removing inner loop. We can utilize the fact that $\text{prefix-sum}[i] = \text{prefix-sum}[i-1] + \text{original array i index value}$.

def calculate_pre_sum_optimized(nums: list[int]):
    prefix_sum = [nums[0]]
    for index in range(1, len(nums)):
        prefix_sum.append(prefix_sum[index - 1] + nums[index])
    
    return prefix_sum

Complexity

$\textbf{Time Complexity} = O(N)$

$\textbf{Space Complexity} = O(N)$

Variations

Prefix sum can have following variations

1. Suffix sum, where we calculate sum from index to end of original array
2. Prefix product, where value at index is product from 0 to index value of original array
3. Prefix min or prefix max, where value at index is min value encountered from 0 to index value of original array
4. 2D prefix sum, where the sum at (i, j) is sum of all elements in rectangle sub-matrix between (0, 0) and (i, j) in original matrix