Subarray Sum Equals K

Keep a running prefix sum and a hash map that records how many times each prefix sum has been seen. For each new element, the sum of any subarray ending here equals the current prefix sum minus some earlier prefix sum, so the number of subarrays ending here with the target sum is exactly the number of times the value current minus target has already appeared.

The seed matters. You start the map with the value zero already present once, because a subarray that starts at index zero has no earlier prefix before it. Without that seed, every subarray that begins at the front of the array would be missed.

The second important detail is the order of the two operations. You look up current minus target before you record the current prefix sum, otherwise a zero target would match the subarray of length zero that you are in the middle of building.

The trade-off to name is that this only works when the array may contain negative numbers, which is exactly why it beats the classic sliding window. A sliding window relies on the sum only moving in one direction as you extend it, and adding a negative number can move it back, so the window has to grow and shrink unpredictably. The prefix sum map has no such assumption and works for any integers, at the cost of linear space. The brute force alternative tries every start and end pair, which is quadratic and easy to write but far too slow on large inputs.

Track the running prefix sum and how many times each prefix sum has been seen. For every element, the number of subarrays that end at that element and sum to the target is the count of earlier prefix sums equal to current minus target.

def subarray_sum(nums, k):
    counts = {0: 1}
    total = 0
    found = 0
    for value in nums:
        total += value
        found += counts.get(total - k, 0)
        counts[total] = counts.get(total, 0) + 1
    return found

print(subarray_sum([1, 1, 1], 2))
print(subarray_sum([1, 2, 3], 3))

The outputs are 2 and 2. In the first array the subarrays that sum to two are the first two elements and the last two, so two. In the second array they are the single element 3 and the pair 1, 2, so two again.

The mechanism is the difference of two prefix sums. The sum of the subarray from index i to j is the prefix sum up to j minus the prefix sum up to i minus one, so a subarray ending at j sums to the target exactly when an earlier prefix sum equals the current one minus the target. Counting how often that earlier value occurred counts all the valid starting points at once, which is why the algorithm is linear rather than quadratic.

The trade-off to name is space. The map holds one entry per distinct prefix sum, so it is linear in the length of the array. The seed of zero with count one is not optional, because it represents the empty prefix that lets a subarray start at index zero. And the lookup has to happen before the update, otherwise a target of zero would count the empty subarray sitting at the current position and over-count the answer by one for every element.

One pass over the array, keeping a running total and a dictionary of the prefix sums seen so far, including the initial zero.

def count_subarrays(nums, target):
    total = 0
    seen = {0: 1}
    count = 0
    for value in nums:
        total += value
        if total - target in seen:
            count += seen[total - target]
        seen[total] = seen.get(total, 0) + 1
    return count

print(count_subarrays([1, 2, 3], 3))
print(count_subarrays([3, 1, 2, 3], 6))

The outputs are 2 and 2. For the first array the valid subarrays are the single 3 and the pair 1, 2. For the second they are the triple 3, 1, 2 and the triple 1, 2, 3, both of which sum to six.

The mechanism is that the dictionary turns the question of how many earlier prefix sums match into a single lookup. The value stored against each prefix sum is a count rather than a flag, because the same prefix sum can occur more than once and each occurrence is a different possible starting point for a subarray.

The trade-off to name is the same one as in the hard answer: linear space for linear time, and the seeded zero entry is what allows subarrays that begin at the first element to be counted. This version writes the lookup as an explicit membership test plus a second read, which is slightly more readable but hashes the key twice; the get with a default of zero does both in one step and is the form to use on a hot path.