Prefix sums: subarrays summing to k

A prefix sum is the running total of a list up to each position. The sum of any slice nums[i:j] is just prefix[j] - prefix[i], so range sums become simple subtraction:

running = 0
for x in nums:
    running += x      # running is the prefix sum so far

This is LeetCode #560. A contiguous subarray ending here sums to k exactly when some earlier prefix equals running - k. So if we remember how many times each prefix value has occurred, we can count matches in one pass.

Your task

Write subarray_sum(nums, k) returning the count of contiguous subarrays whose sum equals k. Numbers may be negative.

subarray_sum([1, 1, 1], 2)    # 2   ([1,1] twice)
subarray_sum([1, 2, 3], 3)    # 2   ([1,2] and [3])
subarray_sum([1, -1, 0], 0)   # 3

Hint: keep a dict of prefix_sum -> how many times seen, seeded with {0: 1} (the empty prefix). For each running sum, add counts[running - k] to your answer, then record running.