Subarray Sum Equals K

Problem

Write a function that returns the total number of contiguous subarrays within a given integer array whose elements sum up to a target K.

Example 1

Input:

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

Output: 4

Explanation: The subarrays that sum to 7 are:

[3, 4], [7], [7, 2, -3, 1], [1, 4, 2]

Example 2

Input:

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

Output: 3

Explanation: The subarrays that sum to 0 are:

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

Explanation

We can solve this question efficiently by leveraging the Prefix-Sum technique.

Let’s start by visualizing the role that the prefix sum plays in this solving this problem. Let’s say I have the following array:

I also have the sum of all the elements starting from the beginning of the array up to index 3 (arr[0] + arr[1] + arr[2] + arr[3]) stored in a variable sum_.

I also have the prefix sums of all the elements up to index 2 stored in an array called prefix:

When looking at these visuals, recall that prefix[i] contains the prefix sum of all elements up to i - 1 in the array.

prefix[i] = arr[0] + arr[1] + ... + arr[i - 1]

Each of the values in the prefix_sum array tells us about the sum of all the different subarrays that end at index 3:

For example, let’s look at prefix[1], which is 3, but we’ll refer to as x for now. From this, we can tell that there exists a subarray ending at index 3 that has a sum of sum_ - x (6 - 3 = 3):

This applies to each value in prefix:

To recap, if we have:

a value sum_ that contains the prefix sum of the array up to index i
another value x that is the prefix sum of the array up to index j (where j comes before i)

then there exists a subarray ending at i with sum sum_ - x.

We’ll use this fact to solve our problem.

Counting Subarrays with Sum K

The question is asking for the number of distinct subarrays that sum to k. So we can turn our prefix array into a dictionary mapping each prefix sum to the number of times it has been seen.

Now, we can use this dictionary to count the number of subarrays that sum to k that end at each index in the array.

Recall from above, we can calculate the sum of any subarray ending at index j by using the difference of prefix sums: if the current prefix sum up to index j is sum_, and there was an earlier prefix sum prev_sum at some index before j, then the subarray sum between them is sum_ - prev_sum.

This means that to find a subarray that sums to k ending at the current index (with prefix sum sum_), we need to find a previous prefix sum equal to sum_ - k in the prefix_counts dictionary. More importantly, the number of subarrays that sum to k that end at index i is equal to prefix_counts[sum_ - k].

Here’s a visual to help us understand this better:

For example, let’s say we have the array below and we are interested in finding the number of subarrays that sum to k = 5. At index 5, we have sum_ = 7

Since we are looking for subarrays that sum to k = 5 and end at index 5, this is equivalent to finding a prefix sum sum_ such that sum_ - k = 7 - 5 = 2. The prefix_counts dictionary tells us that there are 2 of these prefix sums, so there are 2 subarrays that sum to k = 5 and end at index 5:

Solution

At a high level, our solution will iterate over each element in the array and calculate the number of subarrays that sum to k that end at that index using the approach outlined above. We will keep track of the total number of subarrays that sum to k as we iterate over the array, and return it at the end.

initialize a prefix_counts dictionary with a single entry 0: 1. The entry 0: 1 represents that an empty subarray sums to 0, which lets us correctly account for subarrays that start at index 0 and sum to k. The values in the dictionary represent the number of times we have seen a particular prefix sum as we iterate.
initialize sum_ to 0 and count to 0. sum_ represents the prefix sum up to the current index, and count represents the total number of subarrays that sum to k.
iterate over the array, updating sum_ and count as follows: update sum_ by adding the current element to it.
calculate the number of subarrays that sum to k that end at the current index by looking up prefix_counts[sum_ - k] and adding it to count.
update the prefix_counts dictionary by incrementing the count of sum_ by 1.

class Solution:
    def subarraySum(self, nums, k):
        count = 0
        sum_ = 0
        prefix_counts = {0: 1}
        for num in nums:
            sum_ += num
            if sum_ - k in prefix_counts:
                count += prefix_counts[sum_ - k]
            # update the seen prefix sum counts
            prefix_counts[sum_] = prefix_counts.get(sum_, 0) + 1
        return count