Prefix sums and difference arrays¶

~20 min · prerequisite: Two pointers on sorted data

What you will be able to do¶

Build a prefix-sum list and use it to answer the sum of any range in constant time.
Explain why q range queries on a list of length n cost about n + q steps instead of n × q.
Apply many range additions at once with a difference array, then rebuild the values in one pass.

The idea¶

Suppose a judge gives you a list of n numbers and then q questions, each asking for the sum of the numbers from position l to position r. The direct method adds the numbers for each question. A question can cover almost the whole list, so each one costs about n steps. With n and q both equal to 2 × 10^5, that is about 4 × 10^10 steps, which is far too many for a judge.

A prefix sum list stores running totals. Let P[0] = 0, and let P[i] be the sum of the first i values. The list has one more entry than the input. Now look at the difference P[r + 1] - P[l]. It equals the sum of the values at positions l through r, because the first l values cancel out. So each question takes two lookups and one subtraction, which is O(1) steps.

Building the list takes one pass over the input, so it costs O(n) steps. Answering q questions costs O(q) more. The total is O(n + q), which is about 4 × 10^5 steps for the example above.

A difference array works in the other direction. Suppose you must add x to every value from position l to r, and you have many such updates. Doing each one directly costs up to n steps. Instead, add x at index l and subtract x at index r + 1 in an array D of length n + 1. Each update costs O(1). Then take the prefix sums of D once. A prefix sum at index i includes the +x once i reaches l, and it removes that x again once i reaches r + 1. So the result is x exactly for positions l through r. One rebuild costs O(n).

Worked example¶

The function below builds prefix sums with itertools.accumulate, then answers one query.

from itertools import accumulate


def build_prefix(values):
    return [0] + list(accumulate(values))


def range_sum(prefix, l, r):
    return prefix[r + 1] - prefix[l]


values = [3, 1, 4, 1, 5]
prefix = build_prefix(values)
assert prefix == [0, 3, 4, 8, 9, 14]
assert range_sum(prefix, 1, 3) == 6

Line by line, this is what Python does:

accumulate(values) yields the running totals 3, 4, 8, 9, 14. The function list collects them.
[0] + ... puts a 0 at the front, so prefix[i] is the sum of the first i values.
build_prefix returns the list [0, 3, 4, 8, 9, 14]. The assertion checks it.
range_sum(prefix, 1, 3) covers positions 1, 2 and 3, which hold 1, 4 and 1.
It computes prefix[4] - prefix[1], which is 9 - 3 = 6. The assertion checks the answer.

Positions in the question are 0-based, which means the first value has position 0. Keep that convention in mind when a problem numbers things from 1.

Try it¶

The first block applies difference-array updates. The second counts subarrays with a given sum, a problem where prefix sums and a dictionary work together.

from itertools import accumulate


def apply_updates(n, updates):
    diff = [0] * (n + 1)
    for l, r, x in updates:
        diff[l] += x
        diff[r + 1] -= x
    return list(accumulate(diff))[:n]


assert apply_updates(5, [(0, 2, 1), (2, 4, 10)]) == [1, 1, 11, 10, 10]
assert apply_updates(3, []) == [0, 0, 0]

def count_subarrays_with_sum(values, k):
    seen = {0: 1}
    running = 0
    count = 0
    for x in values:
        running += x
        count += seen.get(running - k, 0)
        seen[running] = seen.get(running, 0) + 1
    return count


assert count_subarrays_with_sum([1, 2, 3], 3) == 2
assert count_subarrays_with_sum([1, -1, 1, -1], 0) == 4

Common mistakes¶

Off-by-one in the difference. Using prefix[r] - prefix[l] leaves out the value at position r:
```
prefix = [0, 3, 4, 8, 9, 14]
assert prefix[3] - prefix[1] == 6
```
```
AssertionError
```
The result is 5, which is the sum of positions 1 and 2 only. The correct form is prefix[r + 1] - prefix[l].
Forgetting the leading zero. If you build the list with list(accumulate(values)) alone, then prefix[l - 1] for l = 0 does not fail. Python reads index -1 as the last element, so the answer is silently wrong:
```
from itertools import accumulate

values = [3, 1, 4]
prefix = list(accumulate(values))
assert prefix[1] - prefix[-1] == 4
```
```
AssertionError
```
Negative indices are valid in Python, so no IndexError appears. Add the leading 0 as build_prefix does.
Rebuilding sums inside the query loop. Writing sum(values[l:r + 1]) for every question is correct, but each question costs about n steps. This is the slow method from the idea section, and it is the reason to build prefix sums once.
Mixing 1-based and 0-based positions. If the statement numbers positions from 1, subtract 1 from l and r before you use them. Otherwise the answer may be off by one value, or you may see IndexError: list index out of range when r equals n.

Practice¶

Easy: build the prefix sums of [2, 4, 6] and check the result is [0, 2, 6, 12].
Easy: using the list [1, 2, 3, 4, 5, 6], check that the sum of positions 2 through 4 is 12.
Medium: write answer_queries(values, queries) that takes 0-based inclusive pairs (l, r) and returns the sums in order. Check answer_queries([1, 2, 3, 4], [(0, 0), (1, 3)]) is [1, 9].
Medium: write apply_updates(n, updates) as in the second runnable block. Check apply_updates(4, [(0, 3, 2), (1, 1, 5)]) is [2, 7, 2, 2].
Stretch: write count_subarrays_with_sum and check it against a brute-force count on [2, -2, 2, 2] with k = 2.

Solution

```python from itertools import accumulate

assert [0] + list(accumulate([2, 4, 6])) == [0, 2, 6, 12] ```

python values = [1, 2, 3, 4, 5, 6] assert sum(values[2:5]) == 12

The slice 2:5 holds positions 2, 3 and 4, which are 3, 4 and 5.

```python from itertools import accumulate

def answer_queries(values, queries): prefix = [0] + list(accumulate(values)) return [prefix[r + 1] - prefix[l] for l, r in queries]

assert answer_queries([1, 2, 3, 4], [(0, 0), (1, 3)]) == [1, 9] ```

```python from itertools import accumulate

def apply_updates(n, updates): diff = [0] * (n + 1) for l, r, x in updates: diff[l] += x diff[r + 1] -= x return list(accumulate(diff))[:n]

assert apply_updates(4, [(0, 3, 2), (1, 1, 5)]) == [2, 7, 2, 2] ```

```python def count_subarrays_with_sum(values, k): seen = {0: 1} running = 0 count = 0 for x in values: running += x count += seen.get(running - k, 0) seen[running] = seen.get(running, 0) + 1 return count

values = [2, -2, 2, 2] brute = sum( 1 for i in range(len(values)) for j in range(i + 1, len(values) + 1) if sum(values[i:j]) == 2 ) assert count_subarrays_with_sum(values, 2) == brute ```

The counter seen stores how many earlier prefix sums had each value. A subarray from just after an earlier prefix p to the current prefix s has sum s - p, so the count adds the earlier prefixes equal to s - k.

Recap¶

A prefix sum list turns any range sum into one subtraction, after one O(n) pass.
A difference array turns each range update into two O(1) writes, followed by one prefix-sum pass.
Check the boundaries: the range l to r is prefix[r + 1] - prefix[l], and negative indices in Python do not raise an error.