Medium Prefix Sum + Map Hash Maps
In plain English: Count how many contiguous subarrays add up to exactly k.
If prefix[j] - prefix[i] = k, the subarray i..j sums to k — a hashmap of prefix counts lets you find matching i's in O(1).
Prompt
Count subarrays whose elements sum to exactly k.
Try to write it from scratch before scrolling down.
Solution
def subarray_sum(a, k):
count = 0
prefix = 0
mp = {0: 1}
for x in a:
prefix += x
count += mp.get(prefix - k, 0) # how many earlier prefixes match?
mp[prefix] = mp.get(prefix, 0) + 1
return count
# Test: subarray_sum([1,1,1], 2) == 2