506. Subarray Sum Equals K

We want to find the subarray sums of all arrays.

Imagine the case of [1, 1, 2, 5] where k is 7. We know that [2, 5] is a valid sum, but we want to get the subarray sums in $O(n)$ time, so we’re only allowed a few iterations through the list of nums.

To do this, note that if we iterate through and count the sum:

1, 1, 2, 5
Sums:
1, 2, 4, 9

Note that at the last item, we’ve hit a viable subarray sum. We know this because we have an instance of 2, and 9, where their difference is k (7). In this case, we don’t remember the items that sum up to 7 (in this case, 2 and 5) but note that there must be a subarray from then to here of sum 7.

We apply that to the algorithm — have a dict that counts the sum of the current number, and we increment the amount of times we’ve seen total - k in our map, and use that to increment the total count.

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        sums = defaultdict(int)
        sums[0] = 1
        count = total = 0
        for num in nums:
            total += num
            count += sums[total - k]
            sums[total] += 1
        
        return count