Subarray Sum Equals K

LeetCode 560 · View on LeetCode

Find the number of contiguous subarrays that sum to k.

Brute Force — O(n²)

Check every subarray:

def subarraySum(nums: list[int], k: int) -> int:
    count = 0
    for i in range(len(nums)):
        total = 0
        for j in range(i, len(nums)):
            total += nums[j]
            if total == k:
                count += 1
    return count

Works but too slow for large inputs. Can we do O(n)?

Prefix Sum + Hash Map — O(n)

A subarray nums[i..j] sums to k when two prefix sums differ by k:

prefix_sum_at_j - prefix_sum_at_i = k

Rearranged: at position j, ask "has any earlier prefix sum been equal to current_sum - k?"

nums = [5, 1, 2, 3], k = 3

Prefix sums:  0   5   6   8   11
                              ↑ current sum = 11
                      ↑ earlier sum = 8

11 - 3 = 8  →  "Have I seen 8?"  → YES
               → subarray between them = [3] sums to 3 ✓

The code doesn't do 11 - 8 = 3. It does 11 - 3 = 8 then looks up 8. Same math, just solved for the other side — one O(1) lookup instead of scanning all previous sums.

Why `{0: 1}`?

It handles subarrays starting at index 0. If sum == k, then sum - k = 0, and you need that 0 to exist in the map:

nums = [3], k = 3
sum = 3, need = 3 - 3 = 0

With {0:1}:   found → count = 1 ✓ (the subarray [3] itself)
Without it:   not found → miss it ✗

Why store ALL prefix sums, not just when sum == k?

Because valid subarrays can start anywhere, not just at index 0:

nums = [5, 1, 2, 3], k = 3
No prefix sum ever equals 3! But [1,2] and [3] both sum to 3.

At j=3: sum=11, need=11-3=8, seen[8]=1 → found! subarray [3] ✓
At j=2: sum=8,  need=8-3=5,  seen[5]=1 → found! subarray [1,2] ✓

Every stored prefix sum is a potential LEFT boundary for a future valid subarray.

Full Walkthrough: `nums = [1, 2, 1, 3], k = 3`

seen = {0: 1}  ← "prefix sum was 0 before the array"

j=0: sum = 1
     need = 1 - 3 = -2
     seen[-2]? NO → count stays 0
     Store: seen = {0:1, 1:1}

j=1: sum = 3
     need = 3 - 3 = 0
     seen[0] = 1 → count = 1
     (subarray from start to here: [1,2] sums to 3 ✓)
     Store: seen = {0:1, 1:1, 3:1}

j=2: sum = 4
     need = 4 - 3 = 1
     seen[1] = 1 → count = 2
     (subarray [2,1] sums to 3 ✓ — between where sum was 1 and here)
     Store: seen = {0:1, 1:1, 3:1, 4:1}

j=3: sum = 7
     need = 7 - 3 = 4
     seen[4] = 1 → count = 3
     (subarray [3] sums to 3 ✓ — between where sum was 4 and here)
     Store: seen = {0:1, 1:1, 3:1, 4:1, 7:1}

Answer: 3

Why counts can be > 1

If the same prefix sum appears multiple times, each occurrence is a different start point:

nums = [1, -1, 1], k = 1

j=0: sum=1, need=0, seen[0]=1 → count=1.  seen={0:1, 1:1}
j=1: sum=0, need=-1, no.                   seen={0:2, 1:1}  ← sum 0 seen TWICE now
j=2: sum=1, need=0, seen[0]=2 → count=3.  
     Two start points where sum was 0 → two valid subarrays: [1,-1,1] and [1]

The Code

def subarraySum(nums: list[int], k: int) -> int:
    count = 0
    prefix_sum = 0
    seen = {0: 1}  # prefix sum → how many times seen
    for n in nums:
        prefix_sum += n
        # How many earlier prefix sums are exactly k less than current?
        count += seen.get(prefix_sum - k, 0)
        # Record current prefix sum for future positions
        seen[prefix_sum] = seen.get(prefix_sum, 0) + 1
    return count


if __name__ == '__main__':
    print(subarraySum([1, 2, 1, 3], 3))  # 3
    print(subarraySum([1, 1, 1], 2))     # 2
    print(subarraySum([5, 1, 2, 3], 3))  # 2
    print(subarraySum([1, -1, 0], 0))    # 3