Subarray Sums With a Map

The Idea

Two Sum asks a map for the complement of a value. Do the same with running totals. A subarray ending here sums to k exactly when some earlier prefix total equals running - k. So carry the running total, ask the map how many times that complement has already been seen, add it to the answer, and record the current total for the positions still to come. One pass, and negative numbers do not bother it.

Real-World Example

A bank statement and the question "how many stretches of days net exactly 100?". Running balances answer it: on each day, look up how often the balance once stood 100 lower. Every match marks a stretch that gained precisely that amount.

The Code

def count_subarrays(nums, k):
    seen = {0: 1}          # the empty prefix counts once
    total = running = 0
    for x in nums:
        running += x
        total += seen.get(running - k, 0)      # prefixes that close a k-sum
        seen[running] = seen.get(running, 0) + 1
    return total

print(count_subarrays([1, 2, 3, -2, 2], 3))   # 4

seen = {0: 1}
total = running = 0
for x in nums:
  running += x
  total += seen.get(___, 0)
  seen[running] = seen.get(running, 0) + 1