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.
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.
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.
| Build | Per query | |
|---|---|---|
| Loop each time | none | O(n) |
| Prefix sums | O(n) | O(1) |
Subarrays summing to k
The same array answers a harder question. How many contiguous stretches sum to exactly k?
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 countAt 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:
# add 5 to every element between i and j
diff[i] += 5
diff[j + 1] -= 5Each 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.