Given an array of integers nums and an integer k, return the total number of continuous subarrays whose sum equals to k.
Example 1:
Input: nums = [1,1,1], k = 2 Output: 2
Example 2:
Input: nums = [1,2,3], k = 3 Output: 2
Constraints:
1 <= nums.length <= 2 * 10^4-1000 <= nums[i] <= 1000-10^7 <= k <= 10^7
The description was taken from https://leetcode.com/problems/subarray-sum-equals-k/.
#O(N) Time, O(N) Spaceclass Solution: def subarraySum(self, nums: List[int], k: int) -> int: count, curr_sum = 0, 0 seen = {} for num in nums: curr_sum += num if curr_sum == k: count += 1 if (curr_sum – k) in seen: count += seen[curr_sum - k] if curr_sum not in seen: seen[curr_sum] = 1 else: seen[curr_sum] += 1 return countThis problem is very similar to the Two Sum problem in that we can cache previous values to keep track of what values are in the input.
Instead of caching elements, however, we will be caching sums.
In order to get a subarray that equals k, we may encounter two situations. We have a running sum that started from the beginning of the array that equals k, or we will have a running sum cached that we can exclude from the current sum to also equal k.
This pattern is also known as a Prefix Sum.
Let's start by initializing our count, current sum, and a HashMap for sums that we have seen.
count, curr_sum = 0, 0 seen = {}
Then, we will begin traversing each element from the input array.
We'll start each iteration by incrementing the current sum by the current number.
for num in nums: curr_sum += num
Let's take the input for example, [1,1,1,1], k=2.
If the current sum equals k, then that would be the first situation mentioned previously where we have a running sum that started from the beginning of the array that equals k.
[1,1,1,1], count = 1
if (curr_sum – k) in seen: count += seen[curr_sum - k]
However, if the current sum minus k is in seen, then that would be the second situation mentioned previously where we have a previous running sum cached that we can exclude from the current sum to also equal k.
This falls in line with the x+y=z, x=z-y algorithm that is used in Two-Sum, and it is particularly useful here given that we have negative values because we could utilize those to sums to sum with the current value to get the target sum as well.
[1,1,1,1], count = 2
[1,1,1,1], count = 3
if (curr_sum – k) in seen: count += seen[curr_sum - k]
We cached one as the current sum in the first iteration, so in the third iteration, we can subtract k, two, from the current sum, three, to see if we have any ones that we could exclude from the current sum to get a subarray sum that equals k.
Notice we didn't just increment k by one, but by the number of sums we have seen with that value, this is because we could have multiple sums of the same value that we could combine the current sum with.
At the end of each iteration, we will cache the current sum for later use.
If it's already in seen, we will increment the value. If not, we will initialize it by writing that we have seen it once.
if curr_sum not in seen: seen[curr_sum] = 1 else: seen[curr_sum] += 1
At the end of the traversal, we will return the count of subarrays that sum to k.
return count