Prefix sums and difference arrays

Some problems ask many questions about ranges of one list: the sum from index l to index r, again and again. Some apply many changes to ranges. Doing each one value by value is slow. One pass of preparation makes each one cost O(1).

Prefix sums

Build a list P of length n + 1, where P[i] is the sum of the first i values. So P[0] = 0 and P[i + 1] = P[i] + a[i]. You built the same totals in PYTH-P2; the extra 0 at the front is what makes the next step neat.

P[r + 1] adds up a[0] to a[r]. P[l] adds up a[0] to a[l − 1]. Take one from the other and only a[l] to a[r] is left. Thanks to P[0] = 0, the formula also works for l = 0, with no special case.

from itertools import accumulate
P = [0] + list(accumulate(a))          # O(n), once

def range_sum(l, r):                    # O(1), each time
    return P[r + 1] - P[l]

With q questions, this is O(n + q) in total, instead of O(n × q). The same idea works for anything you can “take away”: counts, or XOR (BITS station). It does not work for the minimum, because a minimum cannot be undone. Range minimums need the structures in the RNGQ station.

Counting with prefix sums

Prefix sums also count ranges. A run a[l..r] adds up to k exactly when P[r + 1] − P[l] = k, which means P[l] = P[r + 1] − k. So go through P once, keeping a dictionary of how many times each earlier value has appeared:

def count_runs(a, k):
    seen = {0: 1}           # P[0] = 0 has appeared once
    s = 0
    count = 0
    for x in a:
        s += x              # s is the next prefix sum
        count += seen.get(s - k, 0)
        seen[s] = seen.get(s, 0) + 1
    return count

This works with negative values too, where a sliding window (lesson 3) fails. It is O(n) on average.

Difference arrays

Now the opposite job. You start with a list of zeros and get many updates of the form “add k to every value from l to r”. At the end you want the final list. Step through the second version of the figure.

Keep a difference list d of length n + 1. For each update, do d[l] += k and d[r + 1] -= k. At the end, the running total of d gives the final values. The +k starts the change at l, and the running total carries it along until the −k cancels it after r.

d[l] += k            # one update: O(1)
d[r + 1] -= k

After the last update, one pass of running totals over d gives the final list, in O(n). Putting the whole method together is ARRS-P3.

Now try ARRS-P3, Passengers on each section.