class Solution:
def subarraySum(self, nums, k):
count = 0
curr_sum = 0
prefix_sums = {0: 1}
for num in nums:
curr_sum += num
count += prefix_sums.get(curr_sum - k, 0)
prefix_sums[curr_sum] = prefix_sums.get(curr_sum, 0) + 1
return count
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
unordered_map<int, int> prefixSums;
prefixSums[0] = 1;
int count = 0, currSum = 0;
for (int num : nums) {
currSum += num;
if (prefixSums.find(currSum - k) != prefixSums.end()) {
count += prefixSums[currSum - k];
}
prefixSums[currSum]++;
}
return count;
}
};
class Solution {
public int subarraySum(int[] nums, int k) {
Map<Integer, Integer> prefixSums = new HashMap<>();
prefixSums.put(0, 1);
int count = 0, currSum = 0;
for (int num : nums) {
currSum += num;
count += prefixSums.getOrDefault(currSum - k, 0);
prefixSums.put(currSum, prefixSums.getOrDefault(currSum, 0) + 1);
}
return count;
}
}
var subarraySum = function(nums, k) {
let count = 0;
let currSum = 0;
const prefixSums = {0: 1};
for (let num of nums) {
currSum += num;
if (prefixSums[currSum - k] !== undefined) {
count += prefixSums[currSum - k];
}
prefixSums[currSum] = (prefixSums[currSum] || 0) + 1;
}
return count;
};
Given an array of integers nums
and an integer k
, find the total number of continuous subarrays whose sum equals to k
.
Key constraints:
nums
can be positive, negative, or zero.nums = [1,2,3], k = 3
2
When first approaching this problem, it's natural to consider all possible subarrays and check if their sum equals k
. This brute-force method involves two nested loops: one to pick a start index and another to pick an end index, summing the elements in between.
However, this approach is inefficient for large arrays because it checks every possible subarray, leading to a time complexity of O(n2).
To optimize, we look for patterns or properties that let us avoid recomputing sums. One key insight is that if we know the sum of elements up to two different indices, we can quickly compute the sum of the subarray between those indices using subtraction. This leads us to the concept of prefix sums and the use of a hash map to store and quickly look up previous sums.
The optimized approach builds on this by maintaining a running sum and counting how often we've seen each possible sum so far, allowing us to find qualifying subarrays in O(n) time.
Let's break down the optimized algorithm step by step:
nums
, we keep a running total called curr_sum
. This represents the sum of all elements from the start up to the current index.
prefix_sums
to store how many times each prefix sum occurs. This helps us quickly check if a subarray sum equals k
.
0
to i
is curr_sum
, and if we've previously seen a prefix sum of curr_sum - k
, then the subarray between those two indices sums to k
.
prefix_sums[0] = 1
to handle the case where a subarray starting from index 0 sums to k
.
nums
:
curr_sum
.curr_sum - k
exists in prefix_sums
, add its count to the answer.curr_sum
in prefix_sums
.
Let's walk through the input nums = [1, 2, 3]
with k = 3
:
Initialization: curr_sum = 0
, prefix_sums = {0:1}
, count = 0
Step 1: num = 1
- curr_sum = 1
- curr_sum - k = 1 - 3 = -2 (not in prefix_sums)
- prefix_sums = {0:1, 1:1}
- count = 0
Step 2: num = 2
- curr_sum = 3
- curr_sum - k = 3 - 3 = 0 (found in prefix_sums, count += 1)
- prefix_sums = {0:1, 1:1, 3:1}
- count = 1
Step 3: num = 3
- curr_sum = 6
- curr_sum - k = 6 - 3 = 3 (found in prefix_sums, count += 1)
- prefix_sums = {0:1, 1:1, 3:1, 6:1}
- count = 2
Result: There are 2 subarrays that sum to 3: [1,2] and [3].
Brute-force approach:
By leveraging prefix sums and a hash map, we efficiently count all subarrays whose sum equals k
in linear time. The key insight is recognizing that the difference between two prefix sums gives the sum of the subarray in between, and that a hash map allows us to quickly find how many times a needed sum has occurred so far. This approach is both elegant and practical, especially for large input arrays.