Subarray Sum Equals K
Problem Description
Given an array of integer nums and an integer k, return the total number of continuous subarrays whose sum equals k.
Examples :
Input: nums = [0, 1, 2, 3, 4, 5, 6], k=9
Output: 2
Explanation: There are 2 subarrays whose sum is 9 i.e. [2, 3, 4] and [4, 5].
Input: nums = [1, -1, 1, -1, 1, -1], k=0
Output: 9
Explanation: There are 9 subarrays whose sum is 0, that are
[1, -1], [1, -1], [1, -1], [-1, 1], [-1, 1], [1, -1, 1, -1], [-1, 1, -1, 1], [1, -1, 1, -1]
and [1, -1, 1, -1, 1, -1].
Constraints:
- 1 ≤ nums.size() ≤ 10⁴
- -10³ ≤ nums[i] ≤ 10³
- -10⁷ ≤ k ≤ 10⁷
Note: There are also some variations of this problem, which are known as count subarray sum equals k. We will also discuss this problem
Brute Force Approach
Let’s begin with a simple brute-force approach. The idea is to iterate over all possible subarrays of the array and check if their sum equals k. If it does, we increment our count.
Iterate through the array:
- We need two nested loops.
- The outer loop starts at each element of the array.
- The inner loop considers all subarrays starting from the current element of the outer loop.
Calculate the sum of each subarray:
- For each subarray defined by the outer and inner loops, we keep a running sum.
- If this running sum equals k, we increment our count.
Here’s a step-by-step breakdown
- Start with the first element.
- Add it to the running sum.
- Check if the running sum is equal to k.
- Move to the next element, add it to the running sum, and check again.
- Continue this process until you’ve considered all subarrays starting from the first element.
- Repeat the process starting from the second element, then the third, and so on.
Explained with an example
Consider the array [0, 1, 2, 3, 4, 5, 6] and k = 9.
Start with the first element (0):
- Subarray [0]: Sum is 0(not equal to 9).
- Subarray [0, 1]: Sum is 1 (not equal to 9).
- Subarray [0, 1, 2]: Sum is 3 (not equal to 9).
- Subarray [0, 1, 2, 3]: Sum is 6 (not equal to 9).
- Subarray [0, 1, 2, 3, 4]: Sum is 10 (not equal to 9).
- Subarray [0, 1, 2, 3, 4, 5]: Sum is 15 (not equal to 9).
- Subarray [0, 1, 2, 3, 4, 5, 6]: Sum is 21 (not equal to 9).
Start with the second element (1):
- Subarray [1]: Sum is 1(not equal to 9).
- Subarray [1, 2]: Sum is 2 (not equal to 9).
- Subarray [1, 2, 3]: Sum is 6(not equal to 9).
- Subarray [1, 2, 3, 4]: Sum is 10(not equal to 9).
- Subarray [1, 2, 3, 4, 5]: Sum is 15(not equal to 9).
- Subarray [1, 2, 3, 4, 5, 6]: Sum is 21(not equal to 9).
Start with the second element (2):
- Subarray [2]: Sum is 2(not equal to 9).
- Subarray [2, 3]: Sum is 5(not equal to 9).
- Subarray [2, 3, 4]: Sum is 9(Equal to 9, count incremented✅).
- Subarray [2, 3, 4, 5]: Sum is 14(not equal to 9).
- Subarray [2, 3, 4, 5, 6]: Sum is 20(not equal to 9).
Continue this process for all starting elements.
Brute Force Code in All Languages
C++ Code
class Solution
{
public:
int subarraySum(vector<int> &nums, int k)
{
// taking the size of the array
int n = nums.size();
// variable to store our count of the subarrays whose sum is k
int ans = 0;
// i is the starting point for each subarray
for (int i = 0; i < n; i++)
{
// Variable to keep the track of subarray sums
int sum = 0;
// j is the ending point for each subarray
for (int j = i; j < n; j++)
{
// add elements to running sum
sum += nums[j];
// if the subarray sum equals to k, increment count
if (sum == k) {
ans++;
}
}
}
// return the answer
return ans;
}
};
Java Code
// Solution class
class Solution {
// Function to count subarrays with sum equal to k
public int subarraySum(int[] nums, int k) {
// Size of the array
int n = nums.length;
// Variable to store count of subarrays whose sum is k
int ans = 0;
// Starting point for each subarray
for (int i = 0; i < n; i++) {
// Variable to keep track of subarray sums
int sum = 0;
// Ending point for each subarray
for (int j = i; j < n; j++) {
// Add elements to running sum
sum += nums[j];
// If the subarray sum equals k, increment count
if (sum == k) {
ans++;
}
}
}
// Return the answer
return ans;
}
}
Python Code
class Solution:
def subarraySum(self, nums, k):
# Get the size of the array
n = len(nums)
# Variable to store the count of subarrays whose sum is equal to k
ans = 0
# Iterate through each starting point for subarrays
for i in range(n):
# Variable to keep track of the current subarray sum
sum = 0
# Iterate through each ending point for subarrays
for j in range(i, n):
# Add current element to the running sum
sum += nums[j]
# If the sum of the current subarray equals k, increment count
if sum == k:
ans += 1
# Return the total count of subarrays with sum equal to k
return ans
Javascript Code
// Solution class with a method to count subarrays with sum equal to k
class Solution {
subarraySum(nums, k) {
// Size of the array
const n = nums.length;
// Count of subarrays whose sum equals k
let ans = 0;
// Iterate through each starting point for subarrays
for (let i = 0; i < n; i++) {
// Track subarray sum
let sum = 0;
// Iterate through each ending point for subarrays
for (let j = i; j < n; j++) {
// Add current element to the sum
sum += nums[j];
// Check if current subarray sum equals k
if (sum === k) {
ans++;
}
}
}
// Return total count of subarrays with sum equal to k
return ans;
}
}
Time Complexity : O(n²)
Why O(n²)?
Explanation: This approach runs two nested loops to find all subarrays, one inside another, where the outer loop runs n times
(i.e. i ranges from 0 to n), and for each iteration of the outer loop, the inner loop(i.e. j ranges from i to n) also runs n-i times.
Analyzing the Loops:
Outer Loop (indexed by i):
It runs from 0 to n-1. So, it iterates n times.
Inner Loop (indexed by j):
For each fixed i, j runs from i to n-1.
When i = 0, j runs n times.
When i = 1, j runs n-1 times.
When i = 2, j runs n-2 times.
…
When i = n-1, j runs 1 time.
Total Iterations:
To find the total number of iterations of the inner loop across all iterations of the outer loop,
we sum up the iterations for each value of i: ∑(i=0 to n−1) (n−i)
This is the sum of an arithmetic series: n+(n−1)+(n−2)+…+1.
If we closely observe and reverse the series
we get 1 + 2 + 3 + … + (n-2) + (n-1) + n.
That is nothing but the sum of n terms.
So the sum of n numbers from 1 to n is (n*(n+1))/2.
Thus, the total number of operations is (n*(n+1))/2.
Simplifying to Big-O Notation:
In Big-O notation, we focus on the term that grows the fastest as n increases, and we ignore constant factors and lower-order terms.
Therefore: O((n*(n+1))/2) ≈ O(n²)
Conclusion: The time complexity of the given code is O(n²).
This is because the nested loops result in a quadratic number of total operations, which dominate the runtime as the input size n becomes large.
Space Complexity : O(1)
Explanation: We need one variable to keep track of the subarray sum, another variable to count the subarrays whose sum equals k, and some i, j pointers to traverse the array, which is constant. Therefore, the space complexity is constant, O(1).
Let’s say you have explained the brute force approach to your interviewer, but they are unsatisfied.
Why is the Interviewer not satisfied?
Because this code will give you a Time Limit Exceed under the given constraints!
Why will it give TLE for n≤10⁵?
If we run it in O(n²) it will take 10¹⁰ operations in the worst case which we can’t run because most competitive programming environments can handle at most 10⁷ to 10⁸ operations within the given time limits.
He would be satisfied with this approach if the constraints for n≤10³ i.e. total at max 10³ * 10³ = 10⁶ operations which is manageable within the typical time constraints of competitive programming.
Now what should come to your mind next to ace it?
If you observe closely, the goal is to find the subarray sum efficiently.
To do this, we need an approach to find the sum of the subarray from the
starting point i to the ending point j in the array nums.
Yes, there is an approach to this. If we store the cumulative sum from index 0 to index i, we can easily know the sum up to the i-th index.
This cumulative sum is often referred to as the prefix sum.
So let's take an array nums = [1, 2, 3, 4, 5, 6]
nums = [1, 2, 3, 4, 5, 6]
prefix(nums) = [1, 3, 6, 10, 15, 21]
Assuming 0-based indexing:
- At index 0: sum = 1
- At index 1: sum = 1 + 2 = 3
- At index 2: sum = 1 + 2 + 3 = 6
- At index 3: sum = 1 + 2 + 3 + 4 = 10
- At index 4: sum = 1 + 2 + 3 + 4 + 5 = 15
So now we know how to find out prefix sum till i-th index, but the problem is we need to find the subarray sum from index i to index j, so how we will find it?
So let's the similar example again nums = [1, 2, 3, 4, 5, 6]
prefix(nums) = [1, 3, 6, 10, 15, 21]
Let’s say we need to find a subarray sum starting at index i
and ending at index j.
Let’s take the example i = 1, j = 4 (assuming 0-based indexing).
So here we know:
- prefix[i-1] = 1
- prefix[j] = 15
To find the sum of the subarray from i to j,
i.e., subarray(nums[i_to_j]), we will simply subtract
prefix[i-1] from prefix[j].
Why i - 1? Because we want to include nums[i] in the sum.
To include nums[i] in the sum, we subtract prefix[i-1] from
prefix[j], which is equivalent to nums[i] + nums[i+1] + ... + nums[j].
- prefix[i-1] = 1
- prefix[j] = 1 + 2 + 3 + 4 + 5 = 15
So, the subarray (nums[i_to_j]) is calculated as:
prefix[j] - prefix[i-1] = (1 + 2 + 3 + 4 + 5) - (1) = (2 + 3 + 4 + 5) = 14.
Let's visualize prefix sum through a pictorial representation :
So in general if we need to find out the subarray sum from index i to j
we will get our answer as prefix[j] — prefix[i-1] (where i-1,j≥0).
What if i=0, i-1 will become negative?
Then the subarray sum will be prefix[j] i.e. a subarray starting with index 0 and ending at index j.
But now we need to find the count of all subarrays whose sum is k.
So how will we be doing it?
Let's observe the prefix sum equation closely.
prefix[j] - prefix[i-1] = k
prefix[i-1] = prefix[j] - k
where i-1<j (i-1 less than j means i-1 lies from 0 to j-1)
So we need to find how many times we have prefix[i-1] = prefix[j]-k in the range from 0 to (j-1), and we will get our job done.
So lets say j=0
we need to find how many times prefix[i-1]=prefix[0]-k occur
in range from 0 -(0-1) (this lie outside the range)
for j=1
we need to find how many times we have prefix[i-1] = prefix[1] - k
in range from 0 to (1-1) = 0.
for j=2
we need to find how many times we have prefix[i-1] = prefix[2] - k
in range from 0 to (2-1) = 1.
for j=3
we need to find how many times we have prefix[i-1] = prefix[3] - k
in range from 0 to (3-1) = 2.
.....so on.
If for every j we run a loop from 0 to j-1, it will again take us to approx. O(n) for getting the count of the number of prefix[i-1] for each prefix[j] in the range 0 to j-1.
So in total, it will take us the same number of operations as
before i.e. n² operation.
To optimize this what if we have a data structure through which we get the count of the number of prefixes from 0 to j-1 in O(1)?
Yes, you’re right! We could use something like a hash table to store the count of each prefix as we go through the array.
What is Hash Table?
A hashtable is a data structure that stores key-value pairs, using a hash function to compute an index for quick access. It offers average O(1) time complexity for operations like insertion, deletion, and lookup.
But wait—our prefix sums could range from -10⁹ to 10⁹, meaning the numbers can be both negative and spread across a huge range. Building a hash table with such a large range would lead to inefficient memory usage.
So, how do we tackle this?
The answer is hash maps. Unlike arrays or hash tables with fixed sizes, a hash map dynamically stores only the elements we encounter, along with their counts. This way, we avoid wasting memory on values that don't exist in the array, making the solution both space-efficient and faster.
What are HashMap's?
Hash maps are a powerful data structure that store key-value pairs, allowing for fast access to values using their keys. They are used to determine where to store each pair, enabling average-case time complexities of O(1) for insertion, search, and deletion.
We can observe things that for each index j we only need the count of prefix[i-1] which is equal to prefix[j]-k where i-1 ranges from 0 to j-1.
So to get the frequency for each prefix[j]-k we need to hash prefix[j] at each iteration for further calculations.
But wait!
We need to do something, for the subarray that starts at index i=0 (beginning of the array nums) and ends at any index j (0≤j≤n-1).
Let's take a small example to understand this :
nums = [0] , k=0
prefix(nums) = [0].
mp = {}
At index j=0,
Check for prefix[j] - k , is prefix[0] - k ie (0-0 = 0) exists in mp ?
No it does not exist.
But we have a subarray right ending at index j=0 with sum=0 right?
So here comes the important of hashing 0 into our map the begininng,
it will ensure that we will consider all the subarray starting at index i=0
and ending at any index j(0<=j<=n-1).
So let redo the same operation after hashing 0 into our map.
mp={{0:1}}
At index j=0,
Check for prefix[j] - k , is prefix[0] - k ie (0-0 = 0) exists in mp?
Yes , it does exists✅.
So the conclusion is we need to initially hash 0 into our map to consider all subarray that starts at the index i = 0 and end at any index j(0≤j≤n-1).
Let's understand the flow of code
Dry run of given example
nums = [1, 2, 3, 4, 5, 6] , k=9.
prefix(nums) = [1, 3, 6, 10, 15, 21]
Hashmap mp={},
Hasing 0 into our map , so mp={{0:1}}
At index j=0,
Check for prefix[j] - k , is prefix[0] - k ie (1-9 = -8) exists in mp ?
No it does not exist.
So count of subarray ending at index j=0 is 0.
Now hash prefix[j] ie 1 into mp
So Hashmap looks like mp={{0:1},{1:1}}
At index j=1,
Check for prefix[j] - k , is prefix[1] - k ie (3-9 = -6) exists in mp ?
No it does not exist.
So count of subarray ending at index j=1 is 0.
Now hash prefix[j] ie 3 into mp
So Hashmap looks like mp={{0:1},{1:1},{3:1}}
At index j=2,
Check for prefix[j] - k , is prefix[2] - k ie (6-9 = -3) exists in mp ?
No it does not exist.
So count of subarray ending at index j=2 is 0.
Now hash prefix[j] ie 6 into mp
So Hashmap looks like mp={{0:1},{1:1}, {3:1}, {6:1}}
At index j=3,
Check for prefix[j] - k , is prefix[3] - k ie (10-9 = 1) exists in mp ?
Yes it does exist ✅.
So we get only 1 subarray starting from 0 to j-1 and ending at index j.
So count of subarray ending at index j=3 is 1.
Now hash prefix[j] ie 10 into mp
So Hashmap looks like mp={{0:1}, {1:1}, {3:1}, {6:1}, {10:1}}
At index j=4,
Check for prefix[j] - k , is prefix[4] - k ie (15-9 = 6) exists in mp ?
Yes it does exist ✅.
So we get only 1 subarray starting from 0 to j-1 and ending at index j.
So count of subarray ending at index j=4 is 1.
Now hash prefix[j] ie 15 into mp
So Hashmap looks like mp={{0:1}, {1:1}, {3:1}, {6:1}, {10:1}, {15:1}}
At index j=5,
Check for prefix[j] - k , is prefix[5] - k ie (21-9 = 13) exists in mp ?
No it does not exist.
So we get only 1 subarray starting from 0 to j-1 and ending at index j.
So count of subarray ending at index j=5 is 0.
Now hash prefix[j] ie 21 into mp
So Hashmap looks like mp={{0:1}, {1:1}, {3:1}, {6:1}, {10:1}, {15:1}, {21:1}}.
So total count of subarray with sum (k = 0) is if we add all the subarray
from index j=0 to j=5 ie 0 + 0 + 0 + 1 + 1 + 0 = 2.
Exit from loop
Let understand it with an visualization
Let's visualize it with an example for array nums = [1, -1, 1, -1, 1, -1].
Intuition
Total number of subarrays with sum equals k in an array =
number of subarrays ending at index 0 with sum k
+ number of subarrays ending at index 1 with sum k
+ ………………………………………….
+ number of subarrays ending at index n-1 with sum k.
So in total if we need to find the count of the subarray that has the sum k we need to find the sum of the count of all subarray ending at index j(0≤j≤n-1) and has the sum k.
Tip : Instead of using a prefix sum array, we could maintain a running sum variable as we have done before in the brute force approach to optimize space.
Optimised Code in All Languages
C++ Code
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
// take the size of the array
int n = nums.size();
// declare an unordered map
unordered_map<int, int> mp;
// to store the number of subarrays having sum as 'k'
int ans = 0;
// to store the prefix for till every ith element'
int sum = 0;
// An empty subarray has a sum of 0, so we hash 0 in the map first
mp[0] = 1;
for(int i = 0; i < n; i++) {
// adding current element to calculate prefix sum till ith index.
sum+=nums[i];
// check if (sum - k) is present in the map
if(mp.count(sum - k)) {
// if yes, add the count of number of prefix[i-1] into the ans
ans += mp[sum - k];
}
// put prefix sum into the map for further calculation
mp[sum]++;
}
// return the answer
return ans;
}
};
Java Code
class Solution {
public int subarraySum(int[] nums, int k) {
// take the size of the array
int n = nums.length;
// Declaring the sum variable to calculate running sum
int sum=0;
// declare a hashmap
HashMap<Integer, Integer> mp = new HashMap<>();
// to store the number of subarrays having sum as 'k'
int ans = 0;
// traverse the prefix array
for (int i = 0; i < n; i++) {
sum+=nums[i];
// if it already equals k, increment count
if (sum == k)
ans++;
// check if (sum - k) is present in the map
if (mp.containsKey(sum - k)) {
// if yes, add to our answer
ans += mp.get(sum - k);
}
// put prefix sum into the map
mp.put(sum, mp.getOrDefault(sum, 0) + 1);
}
// return the answer
return ans;
}
}
Python Code
class Solution:
def subarray_sum(self, nums, k):
# Take the size of the array
n = len(nums)
# Initialize current_sum and a hashmap to store prefix sums
current_sum = 0
mp = {}
# To store the number of subarrays with sum equal to 'k'
ans = 0
# Traverse the array
for i in range(n):
# Update the current sum
current_sum += nums[i]
# Check if current_sum equals k
if current_sum == k:
ans += 1
# Check if (current_sum - k) is in the map
if current_sum - k in mp:
ans += mp[current_sum - k]
# Store current_sum in the map with a count of its occurrences
mp[current_sum] = mp.get(current_sum, 0) + 1
return ans
Javascript Code
class Solution {
subarraySum(nums, k) {
// take the size of the array
const n = nums.length;
// Declaring a sum variable
let sum=0;
// declare a map
const mp = new Map();
// to store the number of subarrays having sum as 'k'
let ans = 0;
// traverse the prefix array
for (let i = 0; i < n; i++) {
sum+=nums[i];
// if it already equals k, increment count
if (sum == k)
ans++;
// check if (sum - k) is present in the map
if (mp.has(sum - k)) {
// if yes, add to our answer
ans += mp.get(sum - k);
}
// put prefix sum into the map
mp.set(sum, (mp.get(sum) || 0) + 1);
}
// return the answer
return ans;
}
}
Time Complexity : O(n)
It is O(n), where n is the size of the array nums.
Explanation: we will be doing our all operations in the single traverse of the array, so for each operation, it takes O(1) time. So in total, we need n*O(1) operations i.e. O(n).
Space Complexity : O(n)
It is O(n), where n is the size of the array nums.
Explanation: For storing the prefix sums in a hash map we need to hash n prefix sums, so in total it takes O(n) space.
Now the interviewer is satisfied with your approach and you have aced the question ✅.
Learning Tip
Since we have solved the Subarray Sum equals k, let’s now tackle the following related problems:
Given an integer array nums and an integer k, return the number of non-empty subarrays that have a sum divisible by k.
Given an integer array nums and an integer k, return true if nums has a good subarray or false otherwise.
A good subarray is a subarray where:
- its length is at least two, and
- the sum of the elements of the subarray is a multiple of k.