-
Notifications
You must be signed in to change notification settings - Fork 0
023 — Prefix Sum
The Arrays and Strings chapter covered the data structure itself — a contiguous, indexable sequence of elements. This chapter is about recognizing when wrapping that array in a running-total transform (the PrefixSum struct's underlying [Int] array) turns repeated range-sum queries from an O(n) re-scan into an O(1) subtraction.
Think of a road trip where your car's odometer only ever ticks upward. You don't reset it at every town — you just note the reading as you pass through each one: 0 miles at the start, 42 at Town A, 97 at Town B, 150 at Town C. If someone later asks "how far is it from Town A to Town C?", you don't re-drive the route — you just subtract: 150 − 42 = 108 miles. The odometer readings are a running total, and any stretch's distance is one subtraction away once you have them. Prefix sums are exactly that odometer: pay the cost of accumulating totals once, up front, and every range query afterward becomes a single subtraction instead of a re-scan.
Reach for prefix sums when you see:
- "Subarray sum equals K" or any variant asking for the count/existence of a contiguous subarray whose sum hits a target — the hallmark move is pairing a running sum with a hashmap of sums seen so far.
-
Repeated range-sum queries on a static (unchanging) array — "answer Q queries of
sum(i, j)" — where recomputing the sum by looping every time would be too slow, but the array itself never changes between queries. - "Equilibrium index," "pivot index," or "find the point where left sum equals right sum" — needs both a running total from the left and the overall total to derive the right side instantly.
- 2D grid versions: "sum of a rectangular submatrix, many queries" — signals a 2D prefix sum table, the natural extension of the same idea.
The unifying tell: the array is queried repeatedly (for ranges, or for "does some earlier sum match") rather than modified, and a naive re-sum on every query would be too slow.
Building a running prefix sum, then answering a range query with one subtraction:
array: [ 2, -1, 3, 4, -2, 1]
index: 0 1 2 3 4 5
prefix[i] = sum of array[0...i-1] (prefix[0] = 0 by convention)
prefix: [0, 2, 1, 4, 8, 6, 7]
index: 0 1 2 3 4 5 6
sum(array[2...4]) = prefix[5] - prefix[2]
= 6 - 1
= 5 ✓ (3 + 4 + -2 = 5)
Same idea powers "subarray sum equals K": walk the array keeping a running sum, and at each step check a hashmap for whether runningSum - K has been seen before — if it has, everything between that earlier point and here sums to exactly K.
target K = 5
array: [1, 2, 3, -2, 5]
running sum: 1, 3, 6, 4, 9
at index 4 (running sum = 9): is (9 - 5 = 4) in the seen-sums map?
yes — it was the running sum right after index 3 →
the subarray from index 4 to 4 sums to 5 ✓ (just [5])
// 1D prefix sum — build once, answer range sums in O(1) each.
struct PrefixSum {
private var prefix: [Int]
init(_ nums: [Int]) {
prefix = [0]
for n in nums {
prefix.append(prefix[prefix.count - 1] + n)
}
}
/// Sum of nums[left...right], inclusive, 0-indexed.
func rangeSum(_ left: Int, _ right: Int) -> Int {
return prefix[right + 1] - prefix[left]
}
}
// The "running sum + hashmap" variant — for counting subarrays that hit a target.
func prefixSumHashmapTemplate(_ nums: [Int], target: Int) -> Int {
var seenSums: [Int: Int] = [0: 1] // running sum 0 seen once, before we start
var runningSum = 0
var count = 0
for n in nums {
runningSum += n
let needed = runningSum - target // 🔧 Fill in: what "complement" are we hunting for?
if let occurrences = seenSums[needed] {
count += occurrences
}
seenSums[runningSum, default: 0] += 1
}
return count // 🔧 Fill in: return whatever the problem actually asks for
}Subarray Sum Equals K — given an integer array nums and an integer k, return the total number of contiguous subarrays whose sum equals k.
func subarraySum(_ nums: [Int], _ k: Int) -> Int {
var seenSums: [Int: Int] = [0: 1]
var runningSum = 0
var count = 0
for n in nums {
runningSum += n
let needed = runningSum - k
if let occurrences = seenSums[needed] {
count += occurrences
}
seenSums[runningSum, default: 0] += 1
}
return count
}
// smoke test
print(subarraySum([1, 1, 1], 2)) // 2
print(subarraySum([1, 2, 3], 3)) // 2 ([1,2] and [3])
print(subarraySum([1, -1, 0], 0)) // 3 ([1,-1], [0], [1,-1,0])
print(subarraySum([3, 4, 7, 2, -3, 1, 4, 2], 7)) // 4This is the hashmap template directly: runningSum accumulates as we scan, and at every index we ask "has runningSum - k shown up before?" — each time it has, every one of those earlier occurrences marks the start of a subarray ending here that sums to exactly k. The seenSums[0] = 1 seed handles the case where the subarray starting at index 0 itself sums to k (its "needed" complement is 0, which must be pre-loaded as having occurred once, representing the empty prefix).
Prefix sum is the odometer trick — accumulate running totals once, and every range becomes a single subtraction instead of a re-drive.
A problem asks for the number of subarrays whose sum is divisible by K (not equal to a fixed target, but divisible by it). The running-sum-plus-hashmap shape still applies — but what do you use as the hashmap's key instead of the raw running sum, and why does that change make "divisible by K" detectable the same way "equals K" was?