Subarray Sum Equals K (LeetCode 560)

The problem

You get an integer array nums (values may be negative) and an integer k. Return how many contiguous subarrays have a sum of exactly k.

Input:  nums = [1, 1, 1], k = 2
Output: 2

The subarrays [1, 1] at indices 0–1 and [1, 1] at indices 1–2 both sum to 2.

How to think about it

Start with brute force. Try every start i, and extend the end j while keeping a running sum. Count every time it equals k. That’s O(n²), which is too slow for 2 × 10⁴ elements in the worst case.

Why a sliding window doesn’t work. A window works when growing it always raises the sum and shrinking it always lowers it. With negative numbers that’s no longer true, so you can’t decide which end to move. You need a different idea.

The key insight: a subarray sum is a difference of two prefix sums. Let prefix[j] be the sum of the first j elements. The subarray from i to j − 1 sums to prefix[j] − prefix[i]. You want that to equal k, which means prefix[i] = prefix[j] − k. So as you walk the array keeping a running sum s, the number of subarrays ending here with sum k is just how many earlier prefix sums equal s − k. A hash map of prefix-sum counts answers that in O(1). This is Two Sum, applied to prefix sums.

Seed the map with {0: 1}. The empty prefix (before any element) has sum 0. Without it, subarrays that start at index 0 would never be counted.

Which pattern is this? Hash map plus prefix sums. The tell is “count subarrays with sum exactly k”, especially when values can be negative.

Approach

1. Create cnt, a map from prefix sum to how many times it has occurred, and set cnt[0] = 1.
2. Keep a running sum s = 0 and an answer res = 0.
3. For each x: add it to s; add cnt[s − k] to res; then record s with cnt[s] += 1.
4. Return res.

Solution

from collections import defaultdict

def subarray_sum(nums, k):
    cnt = defaultdict(int); cnt[0] = 1; s = 0; res = 0   # empty prefix has sum 0
    for x in nums:
        s += x; res += cnt[s - k]; cnt[s] += 1           # earlier prefixes equal to s-k
    return res
# subarray_sum([1,1,1], 2) -> 2

The order inside the loop matters: look up s − k before recording s, so that when k = 0 a prefix isn’t matched with itself.

Dry run

On nums = [1, 1, 1], k = 2 (the map shows only non-zero counts; defaultdict also stores 0 for keys it has looked up):

It returns 2. On [1, 2, 3] with k = 3 it also returns 2 ([1, 2] and [3]), and on [1, −1, 0] with k = 0 it returns 3.

Complexity

• Time: O(n). One pass, with O(1) average map operations per element.
• Space: O(n). Up to n + 1 distinct prefix sums in the map.

Edge cases and pitfalls

• Forgetting cnt[0] = 1 misses every subarray that starts at index 0. On [3] with k = 3, you’d get 0 instead of 1.
• Negative values and zeros: handled naturally; they’re exactly why the sliding window fails.
• k = 0: looking up before recording is what keeps the count right.
• Using a set instead of counts loses multiplicity. With zeros in the array, the same prefix sum repeats and each occurrence is a separate subarray.
• Empty array: returns 0.

Related problems

• Two Sum (LeetCode 1): the same “look up the complement” move, on raw values.
• Contiguous Array (LeetCode 525): map 0 to −1, then find equal prefix sums.
• Subarray Sums Divisible by K (LeetCode 974): the same counting, keyed on prefix sum mod k.