Templates
Template: 13. PREFIX SUM
Mark when done:
13. PREFIX SUM
# =============================================================================
def prefix_sum(arr):
prefix = [0] * (len(arr) + 1)
for i in range(len(arr)):
prefix[i + 1] = prefix[i] + arr[i]
# range sum [i..j] = prefix[j+1] - prefix[i]
return prefix
def subarray_sum_equals_k(arr, k):
"""Count subarrays with sum = k using prefix sum + hashmap."""
count = 0
prefix = 0
seen = {0: 1} # prefix_sum -> count
for num in arr:
prefix += num
if prefix - k in seen:
count += seen[prefix - k]
seen[prefix] = seen.get(prefix, 0) + 1
return count
# =============================================================================