Beautiful Algorithms

Prefix sum

Given an array of integers, return the total number of subarrays whose sum is kk.
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!!