Prefix sums

A step counter keeps how many steps you walked each day. How many steps from day 3 to day 10? sum(steps[3:11]) answers in O(k) for a range of k days. That's fine once. But an app that answers 100,000 such questions about 100,000 days adds up billions of numbers, the same ones again and again.

Prefix sums

Do the adding once. A prefix sum list holds, at index i, the sum of the first i values:

values:     5   2   7   1   4
prefix:  0  5   7  14  15  19

prefix[0] is 0, the sum of no values, so the list is one longer than values. The sum of values[start:end] is then prefix[end] - prefix[start]: prefix[end] adds up everything before end, prefix[start] everything before start, and taking one from the other leaves exactly the values from start to end - 1. One subtraction, O(1), for any range:

def make_prefix(values):
    prefix = [0]
    for value in values:
        prefix.append(prefix[-1] + value)
    return prefix


values = [5, 2, 7, 1, 4]
prefix = make_prefix(values)
print(prefix)
print(prefix[4] - prefix[1], sum(values[1:4]))
print(prefix[5] - prefix[0], sum(values))

[0, 5, 7, 14, 15, 19]
10 10
19 19

Building it is O(n) time and O(n) memory, once. After that, every range sum is O(1). The start and end work like a slice's, so they're easy to remember. The common mistake is to leave out the 0 at the front: then prefix[i] includes values[i], every index is off by one, and a range that starts at 0 needs a special case.

itertools.accumulate

The standard library makes the running totals for you, in C:

>>> from itertools import accumulate
>>> list(accumulate([5, 2, 7, 1, 4]))
[5, 7, 14, 15, 19]
>>> list(accumulate([5, 2, 7, 1, 4], initial=0))
[0, 5, 7, 14, 15, 19]

Without initial=0 you get the totals without the 0 at the front; with it, the prefix list above.

How many stretches add up to k?

A harder question: how many stretches of values in a row (subarrays) add up to exactly k? Trying every start and end is O(n²). Prefix sums turn it into lookups. The stretch from start to end adds up to k when prefix[end] - prefix[start] == k, that is, when prefix[start] == prefix[end] - k. So walk along the list keeping the running total, and at each value count how many earlier prefix sums were total - k. A dict of counts answers that in O(1):

def count_stretches(values, k):
    seen = {0: 1}
    total = 0
    count = 0
    for value in values:
        total += value
        if total - k in seen:
            count += seen[total - k]
        if total not in seen:
            seen[total] = 0
        seen[total] += 1
    return count


print(count_stretches([3, 4, -2, 5, 1, -4, 6], 7))

3

The three stretches are 3, 4, then 4, -2, 5, and then 3, 4, -2, 5, 1, -4. seen maps each prefix sum to how many times it has come up so far, and it starts as {0: 1}: that's the empty prefix before the first value, so stretches that start at index 0 are counted too. Each value costs O(1) on average, so the function is O(n) time and O(n) memory. It works with negative values as well, which the sliding windows of the next step can't handle.

The biggest sum of a stretch

Another classic: which stretch of values in a row has the biggest sum, when values can be negative? The sum of a stretch is prefix[end] - prefix[start], so this is much like best_profit from the first step, run on the prefix sums: remember the smallest prefix sum so far. Only much like, though: best_profit never gives less than 0, as if you could always buy and sell on the same day. On prefix sums that's a stretch of no values at all, so it gives 0 when every value is negative. The biggest sum is usually found another way, known as Kadane's algorithm: keep the best sum of a stretch that ends at the current value. That stretch either starts fresh at the value or adds the value to the best stretch that ended just before it, whichever is bigger. For a list that isn't empty:

def biggest_stretch(values):
    best = values[0]
    ending = 0
    for value in values:
        ending = max(value, ending + value)
        best = max(best, ending)
    return best


print(biggest_stretch([3, -5, 4, -1, 2, 1, -6, 2]))
print(biggest_stretch([-4, -2, -7]))

6
-2

The best stretch is 4, -1, 2, 1. At the 4, the best stretch ending just before it adds up to -2, so starting fresh is better, and the 3 and the -5 are left behind. When every value is negative, the best stretch is the biggest value on its own, -2, which is why best starts at the first value and not at 0. One pass, O(n) time and O(1) memory.