Prefix sums

Prefix sums

Dashboards, banks, and analytics systems answer the same query shape thousands of times a day: what is the total between day lo and day hi?

Re-adding the slice costs O(n) per query, so q queries cost O(n·q). That is the same recompute-from-scratch disease the sliding window cured, except the ranges here are arbitrary rather than a fixed-size window sliding by one.

The fix is a prefix sum array, where prefix[i] is the sum of the first i elements and prefix[0] = 0. It is built once in O(n), and then any range total falls out of one subtraction.

sum of nums[lo..hi] = prefix[hi + 1] − prefix[lo]

The reason is a cancellation. prefix[hi + 1] counts everything up through index hi, and subtracting prefix[lo] removes everything before lo, leaving exactly the range.

So the pattern is build once, then answer every range query in O(1).

The off-by-one is deliberate. Making prefix one cell longer than nums, with that leading 0, lets ranges starting at index 0 use the same subtraction with no special case.

Building the prefix array

The array is one cell longer than nums, and every query is a subtraction.

nums = [3, 1, 4, 1, 5, 9, 2, 6]

prefix = [0]
for n in nums:
    prefix.append(prefix[-1] + n)
print(prefix)

def range_sum(lo, hi):
    return prefix[hi + 1] - prefix[lo]

print(range_sum(2, 5))
print(range_sum(0, 7))

Output

[0, 3, 4, 8, 9, 14, 23, 25, 31]
19
31

The build is one pass, and prefix[-1] + n reuses the running total rather than re-summing, so it is O(n) and not O(n²).

Checking range_sum(2, 5) by hand confirms the rule. The elements are 4 + 1 + 5 + 9 = 19, and the subtraction gives prefix[6] - prefix[2] = 23 - 4 = 19.

The leading 0 earns its place in the second call. range_sum(0, 7) becomes prefix[8] - prefix[0], and without that cell the lo = 0 case would need its own branch.

Both ends of the range are included here, which is a convention to state explicitly. Interview problems differ on it, and mixing the two conventions is where the off-by-one bugs live.

The cost of the structure is O(n) extra memory, and it assumes the data does not change. A single update invalidates every later prefix, which is why mutable data needs a Fenwick tree or segment tree instead.

It costs O(n) once to build and O(1) per query, or about 1,050,000 steps in total.

The build is a single pass of roughly a million additions, and each of the 50,000 queries is one subtraction, so the queries barely register against the build.

The recompute-every-time approach is on the order of 50 billion steps, since an average query might sum hundreds of thousands of entries. That is minutes of work against milliseconds.

The general shape is worth naming, because it recurs. Paying O(n) once to make every later question O(1) is the same buy-structure-first move as sorting in lesson 3-3.

It also tells you when not to build one. For a single query the prefix array is pure overhead, since building it costs the same O(n) as just adding the slice.

Prefix sums plus a dictionary: subarray sums

The famous interview follow-up asks how many contiguous runs sum to exactly k. Checking every (lo, hi) pair is O(n²), and prefix sums turn it into lesson 1-1's pair-sum trick.

A subarray from lo to hi sums to k exactly when prefix[hi + 1] - prefix[lo] = k. Rearranged, that means some earlier prefix value equals current prefix - k.

So the algorithm walks left to right with a running prefix total and a dictionary counting each prefix value seen so far, adding seen[running - k] to the answer at every position.

Counting rather than just recording presence matters, because several earlier positions can share the same prefix value, and each one is a distinct valid subarray.

Seed the dictionary with {0: 1}, the empty prefix, so runs that start at index 0 are counted. One pass gives O(n) time and O(n) space.

Unlike a sliding window, this handles negative numbers, because it never assumes that growing the window grows the sum.

count_subarrays

How many contiguous runs sum to exactly k.

from collections import defaultdict

def count_subarrays(nums, k):
    seen = defaultdict(int)
    seen[0] = 1
    running = 0
    total = 0
    for n in nums:
        running += n
        total += seen[running - k]
        seen[running] += 1
    return total

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

Output

2
2
4

defaultdict(int) returns 0 for unseen keys, so seen[running - k] never raises a KeyError and no membership check is needed.

The order of the three statements is the part to get right. Update running first, then count with seen[running - k], then record seen[running]. Recording before counting would let a zero-length subarray match itself when k = 0.

For [1, 2, 3] with k = 3 the two hits are [1, 2] and [3], and the {0: 1} seed is what catches the first of them, since the run starting at index 0 needs the empty prefix as its left boundary.

The third call is the one a sliding window cannot do. It contains a −3, and the four runs are [3, 4], [7], [7, 2, -3, 1], and [1, 4, 2].

prefix[7] - prefix[3] = 25 - 8 = 17.

The rule is sum of nums[lo..hi] = prefix[hi + 1] - prefix[lo], so with lo = 3 and hi = 6 the indices are 7 and 3.

Reading them off [0, 3, 4, 8, 9, 14, 23, 25, 31] gives prefix[7] = 25 and prefix[3] = 8.

Adding by hand confirms it. nums[3..6] is 1 + 5 + 9 + 2 = 17, which matches.

One subtraction replaced the loop, and that is the entire pattern. The only real skill is keeping the hi + 1 straight, which is why the leading 0 and the inclusive convention are worth stating before writing any code.