Prefix sums and difference arrays

The problem: many range questions

Suppose a shop has daily sales for a year and the manager asks many questions like "what were total sales from day 40 to day 75?". Adding the days each time costs O(n) per question, so q questions cost O(n × q). With a year of data and thousands of questions, that adds up fast.

Prefix sums

A prefix is the start of a list. The prefix sum array holds, at position i, the sum of the first i items. It has n + 1 entries: prefix[0] = 0 (the sum of no items), prefix[1] = nums[0], prefix[2] = nums[0] + nums[1], and so on. Building it takes one pass, O(n) time and O(n) space.

The sum of nums[l] to nums[r] (both included) is then prefix[r + 1] − prefix[l]: the total up to r, minus the total before l. That is O(1) per question.

Example: the sum of nums[1] to nums[3] is prefix[4] − prefix[1] = 9 − 3 = 6, which matches 1 + 4 + 1.

from itertools import accumulate

nums = [3, 1, 4, 1, 5]
prefix = list(accumulate(nums, initial=0))
print(prefix)                    # Output: [0, 3, 4, 8, 9, 14]
print(prefix[3 + 1] - prefix[1]) # Output: 6

Preview: counting subarrays that sum to k

A subarray is a run of neighbouring items. A subarray from l to r sums to k exactly when prefix[r + 1] − prefix[l] = k, that is, when an earlier prefix equals the current prefix minus k. Keep a dictionary that counts how often each prefix has appeared, and you can count all such subarrays in one O(n) pass. Unlike the sliding window of Module 3, this works with negative numbers too.

def count_subarrays(nums, k):
    seen = {0: 1}      # the empty prefix
    running = 0
    count = 0
    for x in nums:
        running += x
        count += seen.get(running - k, 0)
        seen[running] = seen.get(running, 0) + 1
    return count

print(count_subarrays([1, 2, 3], 3))   # Output: 2  ([1, 2] and [3])
print(count_subarrays([1, -1, 0], 0))  # Output: 3

Difference arrays: many range updates

Now flip the problem: many updates, then one read. "Add v to every item from l to r" done directly costs O(r − l + 1) per update. A difference array diff of length n + 1 records only where a change starts and stops: diff[l] += v and diff[r + 1] −= v. After all updates, a running sum over diff gives the final value of every position. Each update is O(1), the final pass is O(n).

Pivot index

The pivot index is where the sum of items to the left equals the sum to the right. With total = sum(nums), walk once and keep left_sum; at index i the right sum is total − left_sum − nums[i]. One pass, O(n) time and O(1) extra space.

Common mistakes

• Off-by-one with the extra 0: the sum of l..r is prefix[r + 1] − prefix[l], not prefix[r] − prefix[l]. Fix: test with l = r and with l = 0.
• Making the difference array length n instead of n + 1, so diff[r + 1] fails for the last index. Fix: allocate n + 1 and ignore the last slot.
• Rebuilding the prefix array inside the query loop. Fix: build it once before answering any query.
• Forgetting seen = {0: 1} in the subarray-sum pattern, which misses subarrays that start at index 0. Fix: the empty prefix counts once.
• Using a sliding window for "sum equals k" when numbers can be negative. Fix: use prefix sums with a dictionary.