backend / dsa / arrays strings / 03_prefix_sums.md

Prefix sums

4 interview angles 3 min read source

Prefix sums

Precompute cumulative sums so any range sum is O(1). The classic space-for-time tradeoff.

When to apply

  • Repeated range-sum queries on a static array
  • “Subarray with sum equals K”
  • “Number of subarrays divisible by K”
  • 2D matrix range queries

Pattern: 1D prefix sum

class PrefixSum:
    def __init__(self, arr: list[int]):
        self.prefix = [0] * (len(arr) + 1)
        for i, x in enumerate(arr):
            self.prefix[i + 1] = self.prefix[i] + x
    
    def range_sum(self, lo: int, hi: int) -> int:
        """Sum of arr[lo..hi] inclusive."""
        return self.prefix[hi + 1] - self.prefix[lo]

prefix[i] = sum of first i elements. Sum of arr[lo..hi] = prefix[hi+1] - prefix[lo]. O(n) preprocessing, O(1) per query.

The off-by-one trick: prefix has length n+1 with prefix[0] = 0. This makes the math work without special casing the boundary.

itertools.accumulate gives you the prefix sum directly:

from itertools import accumulate
prefix = [0] + list(accumulate(arr))

Pattern: subarray sum equals K

def subarray_sum_count(arr: list[int], k: int) -> int:
    """Count subarrays with sum exactly K."""
    count = 0
    prefix = 0
    seen = {0: 1}   # prefix sum 0 occurs once (empty prefix)
    for x in arr:
        prefix += x
        # If prefix - target appeared before at index j, then arr[j+1..i] sums to k
        if prefix - k in seen:
            count += seen[prefix - k]
        seen[prefix] = seen.get(prefix, 0) + 1
    return count

The hash map seen records how many times each prefix sum has occurred. Each query is O(1). Total O(n).

This same idea — “subtract a previous prefix” — solves: longest subarray with sum K, subarray with sum divisible by K (using prefix % k), etc.

Pattern: 2D prefix sum (matrix)

class Matrix2D:
    def __init__(self, mat: list[list[int]]):
        rows, cols = len(mat), len(mat[0])
        self.p = [[0] * (cols + 1) for _ in range(rows + 1)]
        for i in range(rows):
            for j in range(cols):
                self.p[i+1][j+1] = (
                    mat[i][j]
                    + self.p[i][j+1]
                    + self.p[i+1][j]
                    - self.p[i][j]
                )
    
    def query(self, r1: int, c1: int, r2: int, c2: int) -> int:
        """Sum of mat[r1..r2][c1..c2] inclusive."""
        return (
            self.p[r2+1][c2+1]
            - self.p[r1][c2+1]
            - self.p[r2+1][c1]
            + self.p[r1][c1]
        )

Inclusion-exclusion: subtract the two over-counted strips, add back the doubly-subtracted corner.

Problem: equilibrium index

Find an index i where sum(arr[:i]) == sum(arr[i+1:]):

def equilibrium(arr: list[int]) -> int:
    total = sum(arr)
    left = 0
    for i, x in enumerate(arr):
        if left == total - left - x:
            return i
        left += x
    return -1

No explicit prefix array — running prefix sum keeps it O(1) space.

Pitfalls

  • Off-by-one between prefix[i] (sum of first i) vs prefix[i] (sum up to and including index i). Stick to one convention.
  • Forgetting to seed seen[0] = 1 in subarray-sum-equals-K — this represents the empty prefix and is needed when the answer starts at index 0.
  • Integer overflow on big arrays — Python ints are unbounded so this is mostly a non-issue, unlike C++/Java.

Interview angle

  • “What problem do prefix sums solve?” - repeated range-sum queries. Precompute once in O(n), then answer any range in O(1) as prefix[j+1] - prefix[i]. Without it each query is O(n).
  • “Why the extra leading zero?” - a prefix array of length n+1 starting at 0 removes the special case for ranges beginning at index 0, so the formula is uniform. It’s the detail that makes the implementation clean.
  • “Subarray sum equals k?” - prefix sums plus a hash map of counts of previously seen prefix values. For each position, the number of subarrays ending there equals the count of prefix - k already seen. Handles negatives, which a sliding window cannot.
  • “What’s the 2D version?” - a prefix-sum matrix answering any rectangle sum in O(1) via inclusion-exclusion over four corners.
Previous
Sliding window