Prefix Sum
Subarray Sum Equals K
medium
Count: 10
DESCRIPTION (inspired by Leetcode.com)
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.
Solution
Python
class Solution:def subarraySum(self, nums, k):count = 0sum_ = 0prefix_counts = {0: 1}for num in nums:sum_ += numif sum_ - k in prefix_counts:count += prefix_counts[sum_ - k]# update the seen prefix sum countsprefix_counts[sum_] = prefix_counts.get(sum_, 0) + 1return count
Test Your Knowledge
Login to take the complexity quiz and track your progress
Complexity Analysis
Time Complexity: O(N) where N is the length of the input array nums. We iterate over the array once, and for each element, we perform constant time operations.
Space Complexity: O(N) where N is the length of the input array nums. We use a dictionary prefix_counts to store the number of times we have seen each prefix sum, which can have at most N entries.
Mark as read
Unlock Premium Coding Content
Reading Progress
On This Page
Your account is free and you can post anonymously if you choose.