Subarray Sum Equals K

+10

You get an array of integers nums and an integer k. Count the subarrays whose elements add up to exactly k. A subarray is a run of one or more neighbouring elements. Two subarrays count separately when they start or end at different positions, even if they hold the same values. The values may be negative or zero.

Function

subarraySum(nums: integer-array, k: integer) → integer

numsinteger-array: the array of integers, which may hold negative values and zeros
kinteger: the sum a subarray must reach to be counted
Returnsinteger: the number of subarrays whose elements add up to k

Constraints

1 ≤ nums.length ≤ 2 × 10⁴
-1000 ≤ nums[i] ≤ 1000
-10⁷ ≤ k ≤ 10⁷
An array this long has at most 200,010,000 subarrays, so the answer fits in a 32-bit signed integer.

Examples

Input: nums = [3, 4, -7, 1, 3, 3, 1, -4]k = 7
Output: 4
Explanation: Four runs add up to 7: [3, 4], [1, 3, 3], [3, 3, 1] and [3, 4, -7, 1, 3, 3]. In the last one the -7 cancels the 3 and the 4, and the sum climbs back to 7 later, so a run can match even after its sum has passed k.

Input: nums = [1, -1, 0]k = 0
Output: 3
Explanation: Three subarrays add up to 0: [1, -1], [0] and the whole array [1, -1, 0]. The run [-1, 0] adds up to -1, so it does not count.

Input: nums = [2, 2, 2]k = 4
Output: 2
Explanation: The run [2, 2] at indices 0 and 1 and the run [2, 2] at indices 1 and 2 hold the same values but sit at different positions, so both count. The whole array adds up to 6.

+17 hidden tests on Submit

Follow-up

How would you change the solution to return the length of the longest subarray that adds up to k, still in O(n) time?

Hints

Open them one at a time. Each one gives away a little more.

0 of 3 opened

Solution

An array of n numbers has n(n+1)/2 subarrays, about 2 × 10^8 when n = 2 × 10^4, so adding each one up is too slow. The negative values also rule out a sliding window: a window's sum can fall and rise again, so no rule tells you when to shrink it. The idea that cracks the problem is to write every subarray sum as the difference of two prefix sums. Counting subarrays that end at the current element and add up to k then means counting earlier prefix sums equal to the current one minus k, and a hash map answers that in one pass.

Every start with a running total

Correct, but does not finish on the largest tests

Intuition

Every subarray has a first index start and a last index end. If you visit every pair and check its sum, you meet each subarray exactly once, so the count is right.

You do not need a third loop to add each subarray up. Fix start, then move end right one step at a time and add nums[end] to a running total. The total always holds the sum of the elements from start to end, so each subarray costs one addition and one comparison.

Do not stop when the total reaches or passes k. A later negative value can bring it back: in the first example, the total from index 0 goes 3, 7, 0, 1, 4, 7, so that start has a second match at index 5.

The cost is the number of pairs. With n = 2 × 10^4 there are about 2 × 10^8 of them, fine for C, but far too slow for Python, Ruby or R.

Algorithm

Set count to 0.
For each start from 0 to n-1, set total to 0.
For each end from start to n-1, add nums[end] to total.
If total equals k, add 1 to count, and keep going either way.
Return count.

def subarraySum(nums, k):
    n = len(nums)
    count = 0
    for start in range(n):
        total = 0
        # Grow the subarray that begins at start, one element at a time
        for end in range(start, n):
            total += nums[end]
            if total == k:
                count += 1
    return count

Prefix sums with a count map

Intuition

Let prefix[j] be the sum of the first j elements, with prefix[0] = 0 for the empty prefix. The subarray from index i to index j-1 adds up to prefix[j] - prefix[i]. So a subarray that ends at the current element adds up to k exactly when an earlier prefix sum equals the current prefix sum minus k. Each such earlier prefix marks where one matching subarray starts.

Walk the array once. Keep the running prefix sum and a hash map seen from each prefix sum to how many times it has appeared. At each element, first add seen[prefix - k] to the count, then record the current prefix. Looking up before recording keeps a subarray from being empty: with k = 0, recording first would match the current prefix with itself.

Take the first example with k = 7. The prefix sums are 0, 3, 7, 0, 1, 4, 7, 8, 4. When the prefix reaches 7 after index 1, the map holds one 0, which gives [3, 4]. When it reaches 7 again after index 5, the map holds two 0s, the empty prefix and the prefix after the -7, which give [3, 4, -7, 1, 3, 3] and [1, 3, 3] at once. At 8 after index 6, the map holds one 1, which gives [3, 3, 1]. That makes 4.

Starting the map with 0 seen once is what counts the subarrays that begin at index 0. A count map rather than a set matters because the same prefix sum can repeat, and every copy starts a different subarray. Each element costs one lookup and one update, so the time is O(n), and the map holds at most n+1 keys.

Algorithm

Create a map seen with seen[0] = 1, and set prefix and count to 0.
For each element, add it to prefix.
Add seen[prefix - k] to count, reading a missing key as 0.
Add 1 to seen[prefix].
Return count.

def subarraySum(nums, k):
    # seen[p] = how many prefixes so far add up to p; the empty prefix adds up to 0
    seen = {0: 1}
    prefix = 0
    count = 0
    for num in nums:
        prefix += num
        # Each earlier prefix equal to prefix - k starts a subarray that ends here and sums to k
        count += seen.get(prefix - k, 0)
        seen[prefix] = seen.get(prefix, 0) + 1
    return count

Pitfalls and edge cases

Most wrong answers come from treating the input as if every value were positive, or from the order of the two map operations.

A sliding window that shrinks once the sum passes k fails with negative values. On the first example it returns 2 instead of 4: the window keeps its left edge at index 0 until the sum passes 7 at index 6, so it never tries [1, 3, 3] or [3, 3, 1].
Leaving out seen[0] = 1 misses every subarray that starts at index 0. For nums = [5] and k = 5 it returns 0 instead of 1.
Recording the current prefix before the lookup counts empty subarrays when k is 0. For [1, -1, 0] it returns 6 instead of 3.
A set of prefix sums in place of a count map undercounts repeats. For [0, 0, 0] and k = 0 the answer is 6, because each earlier copy of the same prefix sum starts a different subarray.
In the brute force, breaking out of the inner loop when the total goes past k is wrong for the same reason as the sliding window.

Frequently asked questions4

What is the time complexity of Subarray Sum Equals K?

The prefix sum and hash map solution runs in O(n) time and O(n) extra space: one pass, with one lookup and one update per element. Checking every subarray with a running total takes O(n²) time, and adding each subarray up from scratch takes O(n³).

Why does a sliding window not work for Subarray Sum Equals K?

A sliding window relies on the sum rising when the window grows and falling when it shrinks, which holds only when every value is positive. With negative values, a window whose sum is already too large can still become a match after it grows further, so no rule tells you when to move the left edge. If every value were positive, a sliding window would solve it in O(n) time and O(1) space.

Why does the hash map start with 0 mapped to 1?

That entry stands for the empty prefix before the first element, whose sum is 0. A subarray that starts at index 0 adds up to the current prefix sum minus that empty prefix, so without the entry those subarrays are never counted. For nums = [5] and k = 5, the lookup of 5 - 5 = 0 finds that entry and returns 1.

Can Subarray Sum Equals K be solved in O(1) extra space?

Not with the one-pass method. To count the matches that end at an element, you need to know which prefix sums came before it, and there can be up to n+1 different ones. Without the map you fall back to the O(n²) running total. When every value is positive, a sliding window counts the subarrays in O(n) time and O(1) space.

Python

def subarraySum(nums, k):
    # Write code here

Test cases

Case 1

Case 2

Case 3

Input

nums = [3, 4, -7, 1, 3, 3, 1, -4]
k = 7

Expected