Beautiful Algorithms
Prefix sum
Given an array of integers, return the total number of subarrays whose sum is .
LeetCode 560. Subarray Sum Equals K
def solve(nums: List[int], target_sum: int) -> int:
prefixsum2freq = {0: 1} # {prefix_sum: count}
sum = 0
counter = 0
for num in nums:
sum += num # Running sum
overshoot = sum - target_sum # How much we overshot from the target sum
if overshoot in prefixsum2freq: # If True, there was a prefix subarray that can offset the overshoot!!
counter += prefixsum2freq[overshoot]
# Always update the prefix sum at the end!
prefixsum2freq[sum] = prefixsum2freq.get(sum, 0) + 1
return counter
Let's break it down. There are multiple observations
sum
: You notice that this is only tracking the prefix sum, not bothered by anything else- In the same way,
overshoot
is always just showing how much prefix sum is overshot (over the target sum) - A critical observation: the amount we overshot can be offsetted by not including some prefix in the summation!
- This leads to the desire to log what values of prefix-sum we have seen
prefixsum2freq
tells us how many occurrence of a specific prefix sum there were- If an overshot amount matches a prefix-sum we have seen, we can offset!!