Prefix Sums and Difference Arrays in Python: Range Queries and Updates

A weather station keeps daily rainfall in millimetres, and a dashboard asks the same kind of question thousands of times: how much rain fell between day 12 and day 40? Summing that slice each time costs up to n additions per query. If you precompute running totals once, every query becomes a single subtraction. This lesson covers that trick, its mirror image for range updates, and a third use that finds subarrays with a given sum.

Prefix sums: answer any range in one subtraction

Define prefix[i] as the total of the first i values, so prefix[0] is 0. Then the total of days first through last inclusive is prefix[last + 1] - prefix[first]: everything up to and including last, minus everything before first. The extra leading zero is what removes special cases for ranges that start at day 0.

from itertools import accumulate

rain_mm = [0, 4, 12, 0, 0, 7, 22, 3, 0, 1, 15, 6]      # day 0 .. day 11

# prefix[i] = total rain on days 0 .. i-1  (one extra leading 0 keeps the maths clean)
prefix = [0] + list(accumulate(rain_mm))


def rain_between(first_day, last_day):
    """Total rain from first_day to last_day, both inclusive."""
    return prefix[last_day + 1] - prefix[first_day]


print(prefix)
print("days 2-6 :", rain_between(2, 6))
print("days 0-11:", rain_between(0, 11))
print("day 6 only:", rain_between(6, 6))
assert rain_between(2, 6) == sum(rain_mm[2:7])

[0, 0, 4, 16, 16, 16, 23, 45, 48, 48, 49, 64, 70]
days 2-6 : 41
days 0-11: 70
day 6 only: 22

Building the prefix array is one O(n) pass. After that every query is O(1), so q queries cost O(n + q) instead of O(n * q). The trade is memory (one extra list) and one condition: the underlying data must not change between queries. If values are updated constantly, a prefix array goes stale, and structures such as a Fenwick tree are the usual next step.

Difference arrays: update a whole range in O(1)

Now flip the problem. A conference centre has twelve hourly slots and a list of meetings, each needing one room from its first slot to its last. How many rooms are needed at the peak? Adding 1 to every slot of every meeting is a loop inside a loop. A difference array records only the changes: +1 where a booking starts and -1 just after it ends. A running total over the changes then gives the occupancy of every slot.

from itertools import accumulate

HOURS = 12                                   # slots 0..11, i.e. 08:00 to 19:59

# (first_slot, last_slot) inclusive: each is one meeting needing a room
bookings = [(0, 3), (2, 5), (2, 2), (4, 8), (9, 11), (3, 4)]

diff = [0] * (HOURS + 1)                     # one spare cell so last + 1 is always valid
for first, last in bookings:
    diff[first] += 1                         # a room is taken from here on...
    diff[last + 1] -= 1                      # ...and released after `last`

in_use = list(accumulate(diff))[:HOURS]
print("diff  :", diff)
print("in use:", in_use)
print("rooms needed:", max(in_use), "| first peak slot:", in_use.index(max(in_use)))

diff  : [1, 0, 2, 0, 0, -1, -1, 0, 0, 0, 0, 0, -1]
in use: [1, 1, 3, 3, 3, 2, 1, 1, 1, 1, 1, 1]
rooms needed: 3 | first peak slot: 2

Each booking costs two writes no matter how long it lasts, and one pass of accumulate rebuilds the full picture. Prefix sums and difference arrays are inverses: accumulating a difference array gives the values, and differencing the values gives back the difference array. That is why accumulate appears in both.

Prefix sums plus a hash map: exact-sum subarrays

A shop owner keeps a cashbook of daily net movement, positive and negative. How many stretches of consecutive days net to exactly 3? A stretch from day i to j sums to the target when prefix[j + 1] - prefix[i] == target, which means an earlier prefix equal to running - target must exist. Keep a dictionary of how often each prefix value has occurred, and for each day add the number of matching earlier prefixes. This works with negative values, which the sliding window cannot handle.

def count_balanced_stretches(daily_net, target):
    """Count contiguous day ranges whose net cash movement equals `target`."""
    seen = {0: 1}                            # the empty prefix has sum 0, seen once
    running = 0
    found = 0
    for amount in daily_net:
        running += amount
        found += seen.get(running - target, 0)   # earlier prefixes that leave exactly `target`
        seen[running] = seen.get(running, 0) + 1
    return found


cashbook = [3, -1, 2, 1, -3, 4, 0, -2]
print(count_balanced_stretches(cashbook, 3))


def brute(daily_net, target):
    n = len(daily_net)
    return sum(1 for i in range(n) for j in range(i, n) if sum(daily_net[i:j + 1]) == target)


print(brute(cashbook, 3))

4
4

The four stretches are [3], [2, 1], [-1, 2, 1, -3, 4] and [-1, 2, 1, -3, 4, 0]. Note that the last two differ only by a trailing zero, and that the first one starts at day 0, which is exactly why the dictionary begins with {0: 1}. The brute-force version is O(n^3) as written and the fast one is O(n). Writing a slow, obviously correct version beside your clever one is the best way to trust it.

A quick way to decide between these tools is to look at what changes. If the data is fixed and the questions vary, precompute once and answer each question from the table. If the questions are fixed (apply these updates) and you only need the end state, record the changes cheaply and resolve them in one sweep. If you need the best run of a certain shape and values are positive, slide a window. The three techniques share the same skeleton of one pass and a little bookkeeping; they differ in what that bookkeeping remembers.

Common mistakes

• Skipping the leading zero in the prefix array. Symptom: queries starting at index 0 need a special case, and forgetting it raises IndexError or returns a total that is short by the first value.
• Using `prefix[last] - prefix[first]` for an inclusive range. Symptom: every answer is missing the last day's value.
• No spare cell in the difference array. Symptom: IndexError: list index out of range the first time a booking ends in the final slot.
• Forgetting `{0: 1}` in the hash-map version. Symptom: stretches that start at day 0 are not counted; on the cashbook above the result is 3 instead of 4.
• Reusing a prefix array after the data changed. Symptom: queries return old totals with no error. Rebuild it, or use a structure designed for updates.

How this shows up in real projects and interviews

Prefix sums sit behind analytics dashboards that sum a metric over a date range, cumulative histograms and image processing (a 2D version gives the sum of any rectangle). Difference arrays appear in schedulers, booking systems and "apply +x to this range" batch jobs. In interviews, "subarray sum equals k", "range sum query", "corporate flight bookings" and "meeting rooms" are all variants. Good answers name the technique, give both build and query costs, and mention what breaks it (frequent updates). When the whole question is about the minimum or maximum over a window rather than a sum, return to the sliding window lesson.