Subarray Sum Equals K
DESCRIPTION (credit 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.
EXAMPLES
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]
Run your code to see results here
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:
Each of the values in the prefix_sum array tells us about the sum of all the different subarrays tha ends 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 starting from an index i and ending at an index j as sum_ - prefix[j], where sum_ is the prefix sum up to index i.
This means that to find a subarray that sums to k that ends at index i, we need to find a prefix sum sum_ such that sum_ - k is 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 = 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
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.
Loading comments...