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.
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#
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[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:
sum(nums[i..j]) = prefix[j + 1] - prefix[i]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]))8 8
19 19With 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:
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))2
2
6O(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:
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 + 114
14The 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
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]))3
-1
0At 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?