Subarrays summing to k

You'll count subarrays with prefix sums and a hash map. This problem looks almost the same as the shortest subarray reaching a sum: subarrays and a target sum. One sentence differs: here, numbers can be negative. That one sentence breaks the sliding window. Related chapters: prefix sums and hash maps.

1. The problem

You get a list of integers nums, which can include negative numbers and zeros, and an integer k. Return the number of subarrays whose sum is exactly k. A subarray is a run of neighboring numbers. Two runs that look the same but sit at different positions count separately. This is Subarray Sum Equals K (LC 560).

The third example has three answers that overlap and share ends. A method that keeps one window at a time will struggle with it.

2. A correct baseline

For each start, extend a running sum to every end. Unlike the shortest-subarray problem, we can't stop when the sum hits k. A later negative and a later positive number could bring it back to k.

def count_baseline(nums, k):
    count = 0
    for start in range(len(nums)):
        total = 0
        for end in range(start, len(nums)):
            total += nums[end]
            if total == k:
                count += 1
    return count

assert count_baseline([1, 1, 1], 2) == 2
assert count_baseline([1, 2, 3], 3) == 2
assert count_baseline([1, -1, 0], 0) == 3
assert count_baseline([], 0) == 0

This always does n(n + 1) / 2 additions. That's O(n²) time: the work grows with the square of the list's length. With 20,000 numbers, it's about 200 million additions. It uses O(1) extra space.

The bottleneck. For each end, the real question is "how many starts give a sum of exactly k?" The brute force answers it by trying every start again. It asks that question n times, and each answer costs O(n).

3. Why the window from the partner lesson fails

First, see the tempting approach fail. A sliding window for exact sums would extend right, shrink while the sum is over k, and count when the sum equals k:

def count_with_window(nums, k):
    """A tempting adaptation. Correct only for positive values."""
    count = left = window_sum = 0
    for right, value in enumerate(nums):
        window_sum += value
        while window_sum > k and left <= right:
            window_sum -= nums[left]
            left += 1
        if window_sum == k and left <= right:
            count += 1
    return count

assert count_with_window([1, 1, 1], 2) == 2           # agrees here
assert count_with_window([1, -1, 1], 1) == 2          # the true answer is 3
assert count_baseline([1, -1, 1], 1) == 3

At right = 2 in [1, -1, 1], two starts work: index 0 (1 - 1 + 1) and index 2 (1). A window holds one start per end, so it counts at most one of them. With positive numbers, that's fine: at most one start can hit an exact target. With zeros or negatives, several can.

So for each end, we need how many earlier starts give the right sum, without scanning them.

4. The observation

Take [1, 2, 3]. Its running totals are 0, 1, 3, 6, starting from 0 before any number. The sum of [2, 3] is 6 − 1 = 5: the total at the end minus the total before the start. These running totals are called prefix sums.

Write P[i] for the sum of the first i numbers, with P[0] = 0. Then the sum of nums[start .. end] is a difference of two prefix sums:

sum(nums[start .. end]) = P[end + 1] − P[start]

Setting that equal to k and rearranging:

P[start] = P[end + 1] − k

So "how many starts give sum k for this end?" becomes "how many earlier prefix sums equal P[end + 1] − k?" That's a lookup in a hash map of counts, not a scan. Walk left to right, keep a running prefix sum, and count how often each prefix sum has appeared so far.

The rewrite is the key move, not the hash map. A condition on a pair of indexes becomes one where a side depends only on the earlier index. Two Sum uses the same move. Then you can record the earlier side as you go.

5. The solution and a trace

def count_subarrays(nums, k):
    """Number of contiguous subarrays of nums whose sum is exactly k."""
    seen = {0: 1}  # prefix sum -> times seen; the empty prefix has sum 0
    prefix = 0
    count = 0
    for value in nums:
        prefix += value
        count += seen.get(prefix - k, 0)  # earlier starts that complete a sum of k here
        seen[prefix] = seen.get(prefix, 0) + 1
    return count

assert count_subarrays([1, 1, 1], 2) == 2
assert count_subarrays([1, 2, 3], 3) == 2
assert count_subarrays([1, -1, 0], 0) == 3
assert count_subarrays([1, -1, 1], 1) == 3
assert count_subarrays([], 0) == 0

Trace nums = [1, -1, 0] with k = 0. The map is shown after each update.

At the last step, the lookup finds prefix sum 0 twice: before index 0 and before index 2. Those are the subarrays [1, -1, 0] and [0]. One lookup counted both.

6. Why it works

What stays true. Right before the lookup at position end, seen counts the prefix sums P[0] through P[end]. That's true at the start, when only P[0] = 0 exists. Each step adds exactly one new prefix sum, after its lookup.

Each subarray is counted exactly once. A subarray nums[start .. end] sums to k exactly when P[start] == P[end + 1] − k. Since start <= end, that prefix sum is already in seen at step end, so the lookup counts it. Different starts with the same end are counted separately. And each subarray is counted only at its own end, so none is counted twice.

Two details matter. Starting with {0: 1} records P[0], the sum before any number. Without it, you miss subarrays that start at index 0. And the lookup comes before the insert. If you inserted first, then with k = 0 the lookup would find the current prefix sum itself. That counts an empty subarray, which isn't allowed.

def count_insert_first(nums, k):  # deliberately wrong order, for comparison
    seen, prefix, count = {0: 1}, 0, 0
    for value in nums:
        prefix += value
        seen[prefix] = seen.get(prefix, 0) + 1
        count += seen.get(prefix - k, 0)
    return count

assert count_insert_first([5], 0) == 1  # counts an empty subarray; the answer is 0
assert count_subarrays([5], 0) == 0

7. Cost, edge cases, and failure

Cost. One pass, with one hash map lookup and one update per number. That's O(n) time on average, growing in step with the list. The hash map holds up to n + 1 prefix sums, so O(n) space. The answer itself can be as large as n(n + 1) / 2, for example when every number is zero and k = 0. We can count that many subarrays in one pass because we never list them one by one.

Edge cases. An empty list gives 0. k = 0 and runs of zeros are where the order of lookup and insert matters. Python integers never overflow. In other languages, 20,000 numbers of size up to 1,000 sum to at most 20 million, which fits in 32 bits. Check the actual constraints before relying on that.

Where it stops working. A hash map answers "is this exact value here?" If the question becomes "sum at least k" or "sum between two bounds", you need to count earlier prefix sums in a range. A hash map can't do that. Use a structure that keeps order instead, such as a Fenwick tree or merge sort counting, as in Count of Range Sum (LC 327). And if you must list every subarray, the output alone can be O(n²).

import random

rng = random.Random(560)
for _ in range(2000):
    nums = [rng.randint(-3, 3) for _ in range(rng.randint(0, 12))]
    k = rng.randint(-4, 4)
    assert count_subarrays(nums, k) == count_baseline(nums, k), (nums, k)