Prefix Sums with a Hash Map: Subarray Sum and Balance
26 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)Prefix sums answer range questions and hashing answers "seen before?" questions. Together they count blocks with a given total in one pass, and they keep working when numbers are negative.
Counting subarrays the slow way
Subarray is a block of consecutive items of a list, such as items[2:5]. It differs from a subsequence, which may skip items.
A bank wants to know how many runs of consecutive days have a net change of exactly k. Trying every start day and extending the end day works, but takes O(n^2) time.
def count_slow(items, k): count = 0 for i in range(len(items)): total = 0 for j in range(i, len(items)): total += items[j] if total == k: count += 1 return count
print(count_slow([3, 4, 7, 2], 7))Example 1 · Brute force on small lists
def count_slow(items, k): count = 0 for i in range(len(items)): total = 0 for j in range(i, len(items)): total += items[j] if total == k: count += 1 return count
print(count_slow([3, 4, 7, 2], 7))print(count_slow([1, 1, 1], 2))print(count_slow([1, -1, 1, -1], 0))2 2 4
What this code does
- The outer loop picks the start and the inner loop extends the end, keeping a running total.
- Every pair of start and end is visited, so the cost is O(n^2). It handles negative numbers correctly.
A sliding window cannot replace this when items can be negative, because adding an item may lower the sum. We need another idea.
Prefix difference equals k
A block from i to j has sum k when prefix[j + 1] - prefix[i] == k. So at each step, count how many earlier prefixes equal prefix - k. A Counter remembers them.
from collections import Counter
def count_k(items, k): seen = Counter({0: 1}) prefix = count = 0 for x in items: prefix += x count += seen[prefix - k] seen[prefix] += 1 return count
print(count_k([3, 4, 7, 2], 7))Example 2 · Trace of the counting loop
items = [3, 4, 7, 2]k = 7seen = {0: 1}prefix = count = 0for x in items: prefix += x found = seen.get(prefix - k, 0) count += found seen[prefix] = seen.get(prefix, 0) + 1 print(x, prefix, prefix - k, found, count)3 3 -4 0 0 4 7 0 1 1 7 14 7 1 2 2 16 9 0 2
What this code does
- Columns are the item, the prefix, the value looked up, how many earlier prefixes matched, and the running count.
- At prefix 7 the lookup of 0 finds the empty prefix, which gives the block [3, 4]. At prefix 14 the lookup of 7 gives the block [7].
One pass, O(n) time and O(n) space. The lookup comes before the new prefix is stored, so a block is never empty.
Why the map starts with 0: 1
The empty prefix has sum 0, and it happened once. Without it, a block that starts at the first item is never counted. Nothing here needs positive numbers.
from collections import Counter
def count_k(items, k, start=1): seen = Counter({0: start}) prefix = count = 0 for x in items: prefix += x count += seen[prefix - k] seen[prefix] += 1 return count
print(count_k([5], 5), count_k([5], 5, start=0))print(count_k([1, -1, 0], 0))Example 3 · Negatives, zeros and the starting 0
from collections import Counter
def count_k(items, k): seen = Counter({0: 1}) prefix = count = 0 for x in items: prefix += x count += seen[prefix - k] seen[prefix] += 1 return count
print(count_k([1, -1, 0], 0))print(count_k([5], 5))print(count_k([2, -2, 2, -2], 0))3 1 4
What this code does
- The blocks of [1, -1, 0] with sum 0 are [1, -1], [0] and [1, -1, 0].
- The single block [5] is counted because the empty prefix 0 is already in the map.
The first line shows the bug: without the starting 0 the single block [5] is missed. The second line counts blocks with negatives and a zero.
The longest block with a given sum
For the longest block, store the first index where each prefix appeared. An earlier start gives a longer block, so never overwrite a stored index.
def longest_with_sum(items, k): first = {0: -1} prefix = best = 0 for i, x in enumerate(items): prefix += x if prefix - k in first: best = max(best, i - first[prefix - k]) first.setdefault(prefix, i) return best
print(longest_with_sum([1, -1, 5, -2, 3], 3))Example 4 · First index, not last
def longest_with_sum(items, k): first = {0: -1} prefix = best = 0 for i, x in enumerate(items): prefix += x if prefix - k in first: best = max(best, i - first[prefix - k]) first.setdefault(prefix, i) return best
def longest_overwrite(items, k): first = {0: -1} prefix = best = 0 for i, x in enumerate(items): prefix += x if prefix - k in first: best = max(best, i - first[prefix - k]) first[prefix] = i return best
print(longest_with_sum([2, -2, 2, -2], 0))print(longest_overwrite([2, -2, 2, -2], 0))4 2
What this code does
- The first function keeps the earliest index of every prefix and finds the whole list, length 4.
- The second one overwrites the index on every step, so it only sees the latest start and reports a shorter block.
first starts with {0: -1}: the empty prefix ends just before index 0. setdefault keeps the first index of each prefix.
Longest balanced block of 0s and 1s
Balanced subarray is a block of consecutive items in which two kinds of value appear equally often. Scoring one kind +1 and the other -1 turns it into a block with sum 0.
Treat each 1 as +1 and each 0 as -1. A block with equally many of both has sum 0. Seeing the same running balance twice marks such a block.
def longest_balanced(bits): first = {0: -1} balance = best = 0 for i, bit in enumerate(bits): balance += 1 if bit else -1 if balance in first: best = max(best, i - first[balance]) else: first[balance] = i return best
print(longest_balanced([0, 1, 0, 0, 1, 1, 0]))Example 5 · Balanced blocks
def longest_balanced(bits): first = {0: -1} balance = best = 0 for i, bit in enumerate(bits): balance += 1 if bit else -1 if balance in first: best = max(best, i - first[balance]) else: first[balance] = i return best
print(longest_balanced([0, 1, 0, 0, 1, 1, 0]))print(longest_balanced([1, 1, 1]))print(longest_balanced([0, 1]))6 0 2
What this code does
- The balance goes up for 1 and down for 0. A repeated balance means the items in between cancel out.
- When no block is balanced the answer is 0.
Example 6 · Divisible by k, negatives included
from collections import Counter
def count_divisible(items, k): seen = Counter({0: 1}) prefix = count = 0 for x in items: prefix = (prefix + x) % k count += seen[prefix] seen[prefix] += 1 return count
print(count_divisible([4, 5, 0, -2, -3, 1], 5))print(-3 % 5)7 2
What this code does
- Two prefixes with the same remainder mean the items between them add up to a multiple of k.
- In Python,
%with a positive k always gives 0 to k - 1, even for negative numbers, so no extra fix is needed. Some other languages return negative remainders.
Example 7 · Is there a block with sum zero?
from itertools import accumulate
for items in ([4, -3, -1, 5], [1, 2, 3]): prefixes = list(accumulate(items, initial=0)) print(prefixes, len(set(prefixes)) < len(prefixes))[0, 4, 1, 0, 5] True [0, 1, 3, 6] False
What this code does
- Two equal prefix sums mean the items between them add up to zero.
- A set shorter than the list reveals a repeated value, so the check is O(n).
The same idea fits any two kinds of item, such as paid and free orders. It costs O(n) time and O(n) space.
Next step
Check your understanding
12 quick questions