Skip to content

5 of 32 · 5 min

Answering range questions in one step

You are asked for the sum of elements between two indices, and asked it ten thousand times. Each query loops over the range, so each costs O(n) and the whole thing is O(qn).

A single pass beforehand makes every query O(1).

The idea

Build an array where each entry is the total of everything up to that point.

python
nums   = [3, 1, 4, 1, 5, 9, 2]
prefix = [0, 3, 4, 8, 9, 14, 23, 25]

prefix[i] is the sum of the first i elements, with a leading zero so the arithmetic has no special case.

python
def build(nums):
    prefix = [0]
    for n in nums:
        prefix.append(prefix[-1] + n)
    return prefix

def range_sum(prefix, i, j):     # sum of nums[i..j] inclusive
    return prefix[j + 1] - prefix[i]

The sum from i to j is everything up to j, minus everything before i. One subtraction, whatever the size of the range.

BuildPer query
Loop each timenoneO(n)
Prefix sumsO(n)O(1)

Subarrays summing to k

The same array answers a harder question. How many contiguous stretches sum to exactly k?

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

At each position, a stretch ending here sums to k exactly when some earlier prefix equals running - k. Counting how many such prefixes were seen gives the answer in one pass, O(n), where the nested loop is O(n squared).

The difference array, for the opposite problem

When the updates are ranges and the reads come at the end, invert it:

python
# add 5 to every element between i and j
diff[i]     += 5
diff[j + 1] -= 5

Each range update is two operations. A prefix sum over diff at the end reconstructs the final array. A thousand range updates over a million elements becomes O(updates + n) rather than O(updates x n).

Two dimensions

The same construction works on a grid, where the sum of any rectangle is four lookups: the big corner, minus two edges, plus the overlap that was subtracted twice.