Description
Count the non-empty contiguous slices of nums whose elements add up to k. Different start or end positions count separately, even when their values match. Values and k can be negative or zero. Assume 1 to 20,000 elements, each between -1,000 and 1,000, and k between -10,000,000 and 10,000,000.
Solution
def subarray_sum(nums, k):
frequencies = {0: 1}
prefix_sum = 0
count = 0
for value in nums:
prefix_sum += value
count += frequencies.get(prefix_sum - k, 0)
frequencies[prefix_sum] = frequencies.get(prefix_sum, 0) + 1
return count
Examples
Example 1
- Input
[[1,1,1],2]- Output
2
The slices at indices 0..1 and 1..2 both total 2.
Example 2
- Input
[[1,2,3],3]- Output
2
The slices [1, 2] and [3] each total 3.
Example 3
- Input
[[1,-1,0],0]- Output
3
The slices [1, -1], [0], and [1, -1, 0] total zero.
Approach
Maintain a running prefix sum and a map of how often each earlier prefix sum appeared. A slice ending here totals k when its preceding prefix is current_sum - k, so add that prefix's frequency to the answer. Start with frequency 1 for prefix 0 to include slices beginning at index 0. Record the current prefix only after counting matches, which excludes empty slices when k is zero.
Time & space
Expected O(n) time and O(n) auxiliary space, where n is the number of elements in nums. The frequency map holds at most n + 1 distinct prefix sums; hash-map operations take expected O(1) time.