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.

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.

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

Checking up to 2 × 10^8 subarrays is too slow, and negative values rule out a sliding window. A subarray sum is the difference of two prefix sums, so a subarray ending here adds up to k once per earlier prefix sum equal to the current one minus k. A hash map counts those in one pass.

Every start with a running total

Correct, but does not finish on the largest tests

Intuition

Fix the first index start, then move end right and add nums[end] to a running total. The total always holds the sum from start to end, so visiting every pair checks every subarray once, at one addition each.

Do not stop when the total 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, a second match at index 5.

The cost is the number of pairs, about 2 × 10^8 when n = 2 × 10^4.

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 every earlier prefix sum equal to the current prefix minus k marks the start of one matching subarray that ends here. Keep counts, not a set: a prefix sum that repeats starts a different subarray each time.

In the first example the prefix sums are 0, 3, 7, 0, 1, 4, 7, 8, 4. The first 7 finds one 0, giving [3, 4]. The second 7 finds two 0s, giving [3, 4, -7, 1, 3, 3] and [1, 3, 3]. The 8 finds one 1, giving [3, 3, 1].

Look up before you record the current prefix, or with k = 0 the prefix matches itself as an empty subarray. Each element costs one lookup and one update: O(n) time and space.

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

A sliding window that shrinks once the sum passes k fails with negative values: on the first example it returns 2 instead of 4. Breaking out of the brute force's inner loop at that point fails the same way.
Leaving out seen[0] = 1. For nums = [5] and k = 5 it returns 0 instead of 1.
Recording the current prefix before the lookup. With k = 0, [1, -1, 0] returns 6 instead of 3.
A set of prefix sums in place of a count map. For [0, 0, 0] and k = 0 it misses repeats; the answer is 6.

Frequently asked questions3

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.

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

A sliding window needs the sum to rise as the window grows and fall as it shrinks. Negative values break that: a window that is too large can still become a match later, so no rule says when to move the left edge. With only positive values, a window solves 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 0, so without the entry those subarrays are never counted. For nums = [5] and k = 5, the lookup of 5 - 5 = 0 finds it.

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

Test #1

Test #2

Test #3

Arguments

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

Expected Output