Skip to content
Laviyx Learn

Try

DSA course menu

Prefix Sums Explained: Fast Range Sum Queries

Learn prefix sums for O(1) range sum queries, the prefix sum plus hash map pattern for counting subarrays with a given sum, and 2D prefix sums for grids.

Beginner6 min readUpdated
On this page

A prefix sum array stores, at position i, the total of all items before index i. Build it once in O(n), and the sum of any range comes out in O(1).

Building it#

Python
from itertools import accumulate

nums = [3, 1, 4, 1, 5, 9, 2]
prefix = [0]
for x in nums:
    prefix.append(prefix[-1] + x)

print(prefix)
print([0] + list(accumulate(nums)))   # the same, with the standard library
Output
[0, 3, 4, 8, 9, 14, 23, 25]
[0, 3, 4, 8, 9, 14, 23, 25]

The extra 0 at the start makes the formula below work for ranges that begin at index 0.

Range sums in O(1)#

The sum of nums[i] to nums[j] (inclusive) is everything up to j minus everything before i:

Text
sum(nums[i..j]) = prefix[j + 1] - prefix[i]
Python
nums = [3, 1, 4, 1, 5, 9, 2]
prefix = [0]
for x in nums:
    prefix.append(prefix[-1] + x)

def range_sum(i, j):
    return prefix[j + 1] - prefix[i]

print(range_sum(0, 2), sum(nums[0:3]))
print(range_sum(2, 5), sum(nums[2:6]))
Output
8 8
19 19

With many queries this matters: q range questions on a list of n items cost O(n + q) with prefix sums, instead of O(n·q) by summing each range.

Counting subarrays with a given sum#

How many contiguous ranges add up exactly to k? Numbers may be negative.

A range i..j sums to k exactly when prefix[j+1] - prefix[i] == k, so prefix[i] == prefix[j+1] - k. As we move along, a dictionary counts how many earlier prefixes had each value:

Python
def count_subarrays(nums, k):
    seen = {0: 1}          # the empty prefix
    total = count = 0
    for x in nums:
        total += x
        count += seen.get(total - k, 0)
        seen[total] = seen.get(total, 0) + 1
    return count

print(count_subarrays([1, 1, 1], 2))
print(count_subarrays([1, 2, 3], 3))
print(count_subarrays([3, 4, 7, -2, 2, 1, 4, 2], 7))
Output
2
2
6

O(n) time, O(n) space, and unlike the sliding window, it works with negative numbers.

2D prefix sums#

For a grid, P[r][c] holds the sum of the rectangle from the top-left corner to row r − 1, column c − 1. Any rectangle's sum then takes four lookups:

Python
grid = [
    [3, 0, 1, 4],
    [5, 6, 3, 2],
    [1, 2, 0, 1],
]
rows, cols = len(grid), len(grid[0])
P = [[0] * (cols + 1) for _ in range(rows + 1)]
for r in range(rows):
    for c in range(cols):
        P[r + 1][c + 1] = grid[r][c] + P[r][c + 1] + P[r + 1][c] - P[r][c]

def rect_sum(r1, c1, r2, c2):        # inclusive corners
    return P[r2 + 1][c2 + 1] - P[r1][c2 + 1] - P[r2 + 1][c1] + P[r1][c1]

print(rect_sum(0, 0, 1, 1))   # 3 + 0 + 5 + 6
print(rect_sum(1, 1, 2, 3))   # 6 + 3 + 2 + 2 + 0 + 1
Output
14
14

The subtraction removes the strip above and the strip to the left; the top-left corner was removed twice, so it is added back.

Common mistakes#

Practice#

Practice exercise

Write pivot_index(nums): the first index where the sum of items to the left equals the sum of items to the right (not including the item itself), or -1. Use O(1) extra space.

Show hint

Compute the total once. Walk left to right keeping left_sum; the right sum at index i is total - left_sum - nums[i].

Show solution
Python
def pivot_index(nums):
    total = sum(nums)
    left_sum = 0
    for i, x in enumerate(nums):
        if left_sum == total - left_sum - x:
            return i
        left_sum += x
    return -1

print(pivot_index([1, 7, 3, 6, 5, 6]))
print(pivot_index([1, 2, 3]))
print(pivot_index([2, 1, -1]))
Output
3
-1
0

At index 3, the left side is 1 + 7 + 3 = 11 and the right side is 5 + 6 = 11.

Key takeaway#

Next: a structure made of linked nodes, Singly Linked Lists.

Finished this lesson?