Python programs / Data structures / Prefix sums

Python Program for Prefix Sums and Range Sum Queries

Intermediate Data structures Time O(1) Space O(n)

Stores the running totals of a list in one O(n) pass, then answers each range sum query with one subtraction, in O(1).

A prefix sum is the total of a list's items up to a point: the first item, the first two, the first three and so on. Work them all out once, and the total of any stretch of the list, a range sum, takes one subtraction, however long the stretch is.

For [4, 7, 2, 9, 5, 3, 8] the prefix sums are [0, 4, 11, 13, 22, 27, 30, 38]: 0, then 4, 4 + 7 = 11, 11 + 2 = 13, and so on. The total of the items at indexes 2 to 4 (2, 9 and 5) is the total of the first five items minus the total of the first two: 27 - 11 = 16. It's how a program answers many questions about one list that doesn't change, such as the sales between any two dates.

Program

def prefix_sums(items):
    """sums[i] is the total of the first i items."""
    sums = [0]
    for item in items:
        sums.append(sums[-1] + item)
    return sums


def range_sum(sums, start, stop):
    """The total of items[start:stop]."""
    return sums[stop] - sums[start]


numbers = [4, 7, 2, 9, 5, 3, 8]
sums = prefix_sums(numbers)
print("Prefix sums:", sums)
print("Total of numbers[2:5]:", range_sum(sums, 2, 5))
print("Total of numbers[0:4]:", range_sum(sums, 0, 4))

Open in the editor Download prefix_sums.py

Output

Prefix sums: [0, 4, 11, 13, 22, 27, 30, 38]
Total of numbers[2:5]: 16
Total of numbers[0:4]: 22

How it works

  1. sums starts with 0, the total of no items. That makes sums[i] the total of the first i items, and a range that starts at index 0 needs no special case.
  2. Each new total is the one before it plus the next item (sums[-1] is the last total so far), so one pass builds them all. The list it returns has one more entry than items.
  3. The total of items[start:stop] is the total of the first stop items minus the total of the first start items: the items before start are in both totals, so they cancel out. start and stop work as in a slice, so the item at stop isn't included.
  4. The program builds the prefix sums of seven numbers, once, and prints them.
  5. It answers two range queries with one subtraction each: the items at indexes 2 to 4, and the first four items.
i 0 1 2 3 4 5 6 7
numbers[i] 4 7 2 9 5 3 8
sums[i] 0 4 11 13 22 27 30 38
Query Subtraction Answer Check
numbers[2:5] sums[5] - sums[2] = 27 - 11 16 2 + 9 + 5
numbers[0:4] sums[4] - sums[0] = 22 - 0 22 4 + 7 + 2 + 9

Time and space complexity

TimeO(1)
Extra spaceO(n)

Building the prefix sums visits each of the n items once: O(n) time, and O(n) space for the n + 1 totals. After that, every query is one subtraction, O(1), whatever the size of the range. Adding up the range each time instead, with sum(numbers[start:stop]), takes up to O(n) per query, so q queries cost O(n × q) that way against O(n + q) with prefix sums.

Prefix sums suit a list that doesn't change. Changing one item changes every total after it, and redoing them is O(n).

Other ways to write it

With itertools.accumulate()

accumulate() makes the running totals, and initial=0 puts the 0 in front.

from itertools import accumulate

numbers = [4, 7, 2, 9, 5, 3, 8]
sums = list(accumulate(numbers, initial=0))
print(sums, sums[5] - sums[2])

Adding up each range

sum() of a slice gives the same answers with no setup. It's fine for a few queries, but each one goes through every item in its range.

numbers = [4, 7, 2, 9, 5, 3, 8]
print(sum(numbers[2:5]), sum(numbers[0:4]))

Try it yourself

  • Check each answer against sum(numbers[start:stop]), and try a range of one item, range_sum(sums, 3, 4), and an empty one, range_sum(sums, 3, 3).
  • Use range_sum() to find the average of numbers[1:6].
  • Find which 3 items in a row have the largest total, with one range_sum() call for each place they could start.
  • Count the even numbers in any range: build the prefix sums of a list that has 1 for each even number and 0 for each odd one, with a list comprehension.

More data structures programs

Share this program