LeetCode 560: Subarray Sum Equals K

Problem Restatement

We are given an integer array nums and an integer k.

Return the total number of non-empty contiguous subarrays whose sum is exactly equal to k.

A subarray must use consecutive elements.

The problem constraints are 1 <= nums.length <= 2 * 10^4, -1000 <= nums[i] <= 1000, and -10^7 <= k <= 10^7. The array may contain positive numbers, negative numbers, and zero.

Input and Output

Example function shape:

def subarraySum(nums: list[int], k: int) -> int:
    ...

Examples

Example 1:

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

The valid subarrays are:

nums[0:2] = [1, 1]
nums[1:3] = [1, 1]

Both have sum 2.

So the answer is:

2

Example 2:

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

The valid subarrays are:

[1, 2]
[3]

So the answer is:

2

Example 3:

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

The valid subarrays are:

[1, -1]
[0]
[1, -1, 0]

So the answer is:

3

First Thought: Check Every Subarray

A direct solution is to try every possible starting index and every possible ending index.

For each subarray, compute its sum and check whether it equals k.

A basic version looks like this:

class Solution:
    def subarraySum(self, nums: list[int], k: int) -> int:
        ans = 0
        n = len(nums)

        for start in range(n):
            total = 0

            for end in range(start, n):
                total += nums[end]

                if total == k:
                    ans += 1

        return ans

This avoids recomputing each subarray sum from scratch, but it still checks all O(n^2) subarrays.

With n up to 20000, this can be too slow.

Key Insight

Use prefix sums.

Let prefix[i] mean the sum of elements before index i.

Then the sum of a subarray from index j to index i is:

prefix[i + 1] - prefix[j]

We want this to equal k:

prefix[i + 1] - prefix[j] == k

Rearrange it:

prefix[j] == prefix[i + 1] - k

So when we are at the current prefix sum current, we need to know:

current - k

If that prefix sum appeared before, then every previous occurrence gives one valid subarray ending at the current index.

Why a Hash Map Is Needed

We need to count how many times each prefix sum has appeared.

The map stores:

prefix_sum -> frequency

Before reading any element, the prefix sum is 0.

So we start with:

count = {0: 1}

This is important because a subarray starting at index 0 is valid when the current prefix sum itself equals k.

For example:

nums = [1, 2]
k = 3

After reading both numbers, the current prefix sum is 3.

We need:

current - k = 3 - 3 = 0

Since prefix sum 0 appeared once before the array started, [1, 2] is counted.

Algorithm

1. Initialize:

   prefix = 0
   ans = 0
   count = {0: 1}

2. For each number num in nums:

   1. Add it to the running prefix sum.
   2. Compute need = prefix - k.
   3. Add count[need] to the answer.
   4. Record the current prefix sum by incrementing count[prefix].

The order matters.

We count previous prefix sums before inserting the current prefix sum. This prevents counting an empty subarray.

Correctness

At any index, let prefix be the sum of all elements from the start of the array through the current element.

A subarray ending at the current index has sum k exactly when there is an earlier prefix sum equal to:

prefix - k

If that earlier prefix sum occurred c times, then there are exactly c different starting positions for subarrays ending at the current index with sum k.

The algorithm adds exactly this c to the answer.

Then it records the current prefix sum for future subarrays.

Every valid subarray has one ending index. When the algorithm reaches that ending index, it counts the subarray using the prefix sum before the subarray starts. Therefore, every valid subarray is counted once.

The algorithm only counts subarrays whose prefix difference is exactly k, so it never counts an invalid subarray.

Complexity

Let n be the length of nums.

Implementation

from collections import defaultdict

class Solution:
    def subarraySum(self, nums: list[int], k: int) -> int:
        count = defaultdict(int)
        count[0] = 1

        prefix = 0
        ans = 0

        for num in nums:
            prefix += num

            need = prefix - k
            ans += count[need]

            count[prefix] += 1

        return ans

Code Explanation

We create a frequency map for prefix sums:

count = defaultdict(int)
count[0] = 1

The initial 0 prefix represents the sum before the array starts.

Then we track the running sum:

prefix = 0

For each number:

prefix += num

we compute the prefix sum needed to form a subarray sum of k:

need = prefix - k

Then we add the number of previous times that prefix sum appeared:

ans += count[need]

Finally, we record the current prefix sum:

count[prefix] += 1

This makes it available for later subarrays.

Testing

def run_tests():
    s = Solution()

    assert s.subarraySum([1, 1, 1], 2) == 2
    assert s.subarraySum([1, 2, 3], 3) == 2
    assert s.subarraySum([1, -1, 0], 0) == 3
    assert s.subarraySum([0, 0, 0], 0) == 6
    assert s.subarraySum([-1, -1, 1], 0) == 1
    assert s.subarraySum([3, 4, 7, 2, -3, 1, 4, 2], 7) == 4

    print("all tests passed")

run_tests()