Prefix sums in Python

Problem

The dashboard that sums every time

The shop dashboard shows the revenue for a chosen period: “from day 120 to day 340”, “the last 30 days”, “this week against last week”. Managers open it many times a day, and every chart asks about dozens of periods. The code sums the days of every period again:

def range_sum(revenue, a, b):
    total = 0
    for day in range(a, b + 1):
        total += revenue[day]
    return total

A thousand days of revenue and a thousand queries:

revenue = revenue_of(1000)
queries = queries_of(1000)
total = 0
additions = 0
for a, b in queries:
    total += range_sum(revenue, a, b)
    additions += b - a + 1
print(f"queries: {len(queries)}")
print(f"sum of the answers: {total}")
print(f"additions: {additions}")

queries: 1000
sum of the answers: 68412493
additions: 248794

248,794 additions for a thousand answers. The days do not change between two queries, yet the code adds them up again and again. The longer the periods and the more queries there are, the slower the dashboard gets.

Where the days and the queries come from

The revenue and the queries follow a formula, so that the programs in all four languages and the comparison in the browser below build exactly the same ones:

def revenue_of(n):
    return [50 + (i * 37) % 451 for i in range(n)]


def queries_of(n):
    queries = []
    for k in range(n):
        a = (k * 7919) % n
        queries.append((a, a + (k * 4231) % (n - a)))
    return queries

The idea

The sum of everything to the left

Compute once, for every day, the revenue from the start up to that day. This is a prefix sum: prefix[k] is the sum of the first k days.

prefix[0] = 0: the sum of no days at all;
prefix[k + 1] = prefix[k] + revenue[k]: every place is the previous one plus one day.

The array has one more place than there are days. After that, every period is the difference of two numbers. The revenue from day a to day b is days 0 to b minus days 0 to a − 1:

prefix[b + 1] − prefix[a]

Building it is n additions, once. Every query after that is two reads and one subtraction, whether the period is two days or two years: O(1) per query instead of O(n).

See it

The prefixes, day by day

def build_prefix(revenue):
    prefix = [0] * (len(revenue) + 1)
    for i, amount in enumerate(revenue):
        prefix[i + 1] = prefix[i] + amount
    return prefix


def range_sum(prefix, a, b):
    return prefix[b + 1] - prefix[a]

revenue = [120, 80, 150, 90, 200, 60, 110, 140]
prefix = build_prefix(revenue)
print("prefix:", prefix)
print(f"days 2–5: {range_sum(prefix, 2, 5)}")

prefix: [0, 120, 200, 350, 440, 640, 700, 810, 950]
days 2–5: 500

def build_prefix(revenue):
    prefix = [0] * (len(revenue) + 1)
    for i, amount in enumerate(revenue):
        prefix[i + 1] = prefix[i] + amount
    return prefix


def range_sum(prefix, a, b):
    return prefix[b + 1] - prefix[a]

Step by step, in words

8 days of revenue. prefix has 9 places: prefix[0] = 0 is the sum of no days at all.
prefix[1] = prefix[0] + revenue[0] = 0 + 120 = 120: days 0–0.
prefix[2] = prefix[1] + revenue[1] = 120 + 80 = 200: days 0–1.
prefix[3] = prefix[2] + revenue[2] = 200 + 150 = 350: days 0–2.
prefix[4] = prefix[3] + revenue[3] = 350 + 90 = 440: days 0–3.
prefix[5] = prefix[4] + revenue[4] = 440 + 200 = 640: days 0–4.
prefix[6] = prefix[5] + revenue[5] = 640 + 60 = 700: days 0–5.
prefix[7] = prefix[6] + revenue[6] = 700 + 110 = 810: days 0–6.
prefix[8] = prefix[7] + revenue[7] = 810 + 140 = 950: days 0–7.
8 additions, once. Now the query: the revenue from day 2 to day 5. prefix[6] = €700 is days 0–5.
prefix[2] = €200 is days 0–1, the part before the period. Take it away.
700 − 200 = €500: days 2–5 with two reads and one subtraction, however long the period is.
Building costs n additions once; every query after that costs one subtraction.

In the scene prefix[k] sits on the boundary between day k − 1 and day k. That shows it is the sum of everything to the left of the boundary, and a period is the difference between two boundaries.

More precisely, if you are curious

What the languages give you.

itertools.accumulate gives the prefixes as an iterator. With initial=0 it starts with the zero in front, exactly like build_prefix:

from itertools import accumulate

revenue = [120, 80, 150, 90]
print(list(accumulate(revenue, initial=0)))

[0, 120, 200, 350, 440]

Overflow. In Python an integer has no upper limit, so even the largest prefix is exact. In Java and C# the array is long, because the last prefix (years of revenue, in cents) goes past the limit of int.

Try it

Predict prefix

Three days of revenue: €30, €50 and €20.

revenue = [30, 50, 20]
prefix = build_prefix(revenue)
print("prefix:", prefix)

Answer

prefix: [0, 30, 80, 100]

prefix: [30, 80, 100]

prefix: [0, 30, 50, 20]

Try it

Fill in

The revenue from day a to day b inclusive:

return prefix[] - prefix[a]

The idea

The zero in front

Why does the array have one more place? It can do without: prefix[k] could be the sum of days 0 to k, inclusive. Then the period from a to b is prefix[b] − prefix[a − 1]. But for a = 0 there is no place a − 1, and every query needs a check “if a is 0”.

The zero in front removes that check. The price is one place in the array, and the trap is that the indexes shift by one: prefix[k] is the sum of the days before k, not up to k. That is why an off-by-one error is the most common mistake here, as with binary search.

Try it

Find the bug

Someone wrote the prefixes without the zero in front. The query for days 2–5 should give €500:

revenue = [120, 80, 150, 90, 200, 60, 110, 140]
prefix = build_prefix(revenue)
print(f"days 2–5: {range_sum(prefix, 2, 5)}")

days 2–5: 350

Pick the line with the bug.

See it

Measure: re-summing against the prefixes

The same thousand queries, with the prefixes:

revenue = revenue_of(1000)
queries = queries_of(1000)
prefix = build_prefix(revenue)
operations = len(revenue)
total = 0
for a, b in queries:
    total += range_sum(prefix, a, b)
    operations += 1
print(f"queries: {len(queries)}")
print(f"sum of the answers: {total}")
print(f"additions and subtractions: {operations}")

queries: 1000
sum of the answers: 68412493
additions and subtractions: 2000

The sum of the answers is the same, and the work is 2,000 operations instead of 248,794: a thousand additions for prefix and one subtraction per query. Run both in the browser. n is both the number of days and the number of queries:

Re-summing for every query O(n²)
Prefix sums O(n)

Algorithm	n = 10	n = 1,000	n = 100,000
Re-summing for every query	10	248,794	≈ 2,487,940,000
Prefix sums	20	2,000	200,000

The code runs in your browser, in JavaScript. The number of operations is the same in every language, because the algorithm is the same. The time depends on the machine, so it is only for comparison. “≈” is an estimate from the curve: we do not run an n that big, so the page does not freeze.

def build_prefix(revenue):
    prefix = [0] * (len(revenue) + 1)
    for i, amount in enumerate(revenue):
        prefix[i + 1] = prefix[i] + amount
    return prefix


def range_sum(prefix, a, b):
    return prefix[b + 1] - prefix[a]

Re-summing is O(n) per query, O(n²) for n queries. With the prefixes it is O(n) once and O(1) per query. The price is memory: one array with n + 1 numbers.

Two-dimensional prefix sums: orders by day and hour

The heat map of orders is a table: the rows are days of the week, the columns are hours. The question is “how many orders are there on Tuesday and Wednesday from 10:00 to 13:00”, that is, the sum inside a rectangle.

p[r][c] is the sum of the rectangle from the top left corner to row r − 1 and column c − 1, again with zeros in front (row 0 and column 0). Every place is the cell plus the rectangle above and the one to the left. The corner between them was counted twice, so it is taken away once:

def build_prefix_2d(orders):
    rows, cols = len(orders), len(orders[0])
    p = [[0] * (cols + 1) for _ in range(rows + 1)]
    for r in range(rows):
        for c in range(cols):
            p[r + 1][c + 1] = orders[r][c] + p[r][c + 1] + p[r + 1][c] - p[r][c]
    return p


def rect(p, r1, c1, r2, c2):
    return p[r2 + 1][c2 + 1] - p[r1][c2 + 1] - p[r2 + 1][c1] + p[r1][c1]

orders = [
    [3, 8, 12, 9, 5, 4],
    [2, 7, 14, 11, 6, 3],
    [4, 9, 15, 10, 7, 5],
    [1, 6, 11, 8, 4, 2],
]
p = build_prefix_2d(orders)
print(f"orders Tue–Wed from 10:00 to 13:00: {rect(p, 1, 1, 2, 3)}")

orders Tue–Wed from 10:00 to 13:00: 66

The query is the same idea: the whole rectangle up to the bottom right corner, minus the strip above, minus the strip to the left, plus the corner that was taken away twice. Four reads for a rectangle of any size.

The idea

How many periods make exactly €50

The shop has days with a loss: returned orders are taken off the revenue. The question is how many periods of consecutive days have a profit of exactly €50.

Every period from day i to day j is the difference of two prefixes: prefix[j + 1] − prefix[i]. It is 50 exactly when prefix[i] = prefix[j + 1] − 50. So you walk through the days, keep the current prefix and ask: how many times so far have I seen a prefix equal to the current one minus 50? The answer is kept in a map from prefix to how many times it was seen:

def count_periods(profit, target):
    seen = {0: 1}
    prefix = 0
    count = 0
    for p in profit:
        prefix += p
        count += seen.get(prefix - target, 0)
        seen[prefix] = seen.get(prefix, 0) + 1
    return count

profit = [40, -10, 30, 20, -30, 50, -20, 30, 0, 20]
print(f"periods with a profit of €50: {count_periods(profit, 50)}")

periods with a profit of €50: 8

One pass over the days and one lookup in a hash table per day: O(n). Two nested loops over all periods would be O(n²). The map here is a counter: how many times each value was seen. That is a technique of its own, and it is the subject of the next lesson.

More precisely, if you are curious

seen starts with 0 → 1, because the prefix before the first day is 0. Without it the periods that start on day 0 are not counted.

If every day made a profit, a sliding window would also work: grow it on the right while the sum is under 50, shrink it on the left once it goes over. With negative numbers that fails: shrinking can increase the sum, and the window cannot tell which way to go. Prefixes with a map work with any numbers.

Try it

Predict the count

Four days of profit: €20, €30, −€30 and €30. How many periods of consecutive days make exactly €50?

profit = [20, 30, -30, 30]
print(f"periods with a profit of €50: {count_periods(profit, 50)}")

Answer

periods with a profit of €50: 3

periods with a profit of €50: 1

periods with a profit of €50: 2

Exercise

The warehouse, day by day

Solve the task in your own editor.

Delivered, sold, in stock. The warehouse starts with 40 units of one product. For ten days you have two arrays:

delivered: 0, 25, 0, 30, 10, 0, 40, 0, 15, 0;
sold: 12, 8, 15, 9, 20, 11, 7, 18, 10, 14.

Build the prefixes of both arrays. For the periods days 2–5, 0–9 and 6–8 print how much was delivered, how much was sold, and the stock at the end of the last day. Every answer must come from a few reads of the prefixes, without a loop over the days.

days 2–5: delivered 40, sold 55, stock at the end 30
days 0–9: delivered 120, sold 124, stock at the end 36
days 6–8: delivered 55, sold 35, stock at the end 50

More tasks

The average order in a period. Besides the revenue you also have the number of orders per day. An average order value does not add up, but the two sums do. How do you answer “the average order from day a to day b”?
The best week. With the prefixes, find the seven consecutive days with the highest revenue. How many operations does that take for 365 days?
The heat map. Add Friday and Saturday to the two-dimensional sums and find the two-hour window with the most orders in the whole week.

Test

Check yourself

1. The revenue for four days is €10, €20, €30 and €40. What is prefix?[0, 10, 30, 60, 100][0, 10, 20, 30, 40][10, 30, 60, 100]

2. Which one gives the revenue from day a to day b inclusive?prefix[b + 1] − prefix[a]prefix[b + 1] − prefix[a + 1]prefix[b] − prefix[a]

3. Yesterday’s revenue is corrected because of a returned order. What happens to prefix?Only one place changesNothing: prefix does not depend on yesterdayEvery place after that day changes: the correction costs O(n)

4. In countPeriods the map seen starts with 0 → 1. Why?So that it is not emptyBecause of the periods that start on day 0: for them we look for the prefix 0 from before the first dayIt does not matter, it only speeds things up

Prefix sums

Cookies