Count Number of Subarrays with Sum K

Introduction

Problems of arrays of finding subarrays and subsequences satisfying certain conditions are frequently asked in Amazon, Microsoft, and Google.
Today this article will discuss one of those problems: Subarray Sum, where we have to count the number of subarrays with sum k.

Let’s understand the problem statement more precisely in the next section.

Also, see Data Structures.

Problem Statement

Given an integer array and an integer k. Find the total number of contiguous subarrays of a given array whose sum of elements is equal to k.

Example:

Suppose given ‘arr’ is [3, 1, 2, 4] and ‘k’ is 6 then

The count of the subarrays in ‘arr’ that sums up to 6 is 2, i.e. subarrays ‘arr[0...2]’ and ‘arr[2...3]’.

Approach 1: Brute Force

A simple solution could be to traverse all the subarrays and calculate their sum. If the sum is equal to the given required sum, then increment the count of subarrays. Finally, return the count of such subarrays.

The algorithm to calculate the required will be:-

Initialize a variable ‘res’ as 0 to store the count of all the subarrays which sum up to the given value ‘k’.
Now, traverse the given array by iterating from 0 to ‘N’, which is the total size of the array and traverse again from the current index up to the end of the array (basically getting the subarrays).
Keep check If there exists a subarray from index up to ‘N’ having the sum equal to ‘k’ and keep incrementing the value of ‘res’ if it does otherwise, keep moving ahead and checking for all the other subarrays present in ‘arr’.
Finally, return the result ‘res’.

C++ Code

/*
    Time Complexity: O(N^2)
    Space Complexity: O(1)

    Where 'N' is the number of elements in the array/list 'arr'.
*/
int findAllSubarraysWithGivenSum(vector < int > & arr, int k) {
    int n = arr.size();
    int res = 0;

    // Calculate all subarrays.
    for (int i = 0; i < n; i++) {
        int summ = 0;
        for (int j = i; j < n; j++) {
            // Calculate required sum.
            summ += arr[j];

            // Check if the sum is equal to the required sum.
            if (summ == k) {
                res += 1;
            }
        }
    }
    return res;
}

// main program to implement above functions
int main()
{
  	vector <int> arr = {5,0,5,10,3,2,-15,4};
  	int k = 5;
  	
  	cout<<"Number of subarrays with sum equal to "<<k<<" = "<<findAllSubarraysWithGivenSum(arr,k)<<endl;
  	
 	return 0;
}

Java code

import java.io.*;
class Solution {
    
	public static void main (String[] args) {
    	int[] arr = {5,0,5,10,3,2,-15,4};
    	
    	//k is the given target sum
    	int k = 5;
    	
    	//returns the number of subarrays with sum k
    	int res = subarraysWithGivenSum(arr, k);
    	
    	//print the result
		System.out.println(res);
	}

	//function to find the number of
	//subarrays with given sum
	public static int subarraysWithGivenSum(int[] arr,int k){
    	int len = arr.length;
    	
    	//stores the count of subarrays
    	int ans = 0;
    	
    	//calculate all subarrays
    	for(int i = 0; i < len; i++){
        	int currSum = 0;
        	for(int j = i; j < len; j++){
            	//calculate required sum
            	currSum += arr[j];
            	
            	//check if sum is equal to required sum
            	if(currSum == k){
                	ans += 1;
            	}
        	}
    	}
    	//return the ans
    	return ans;
	}
}

Python Code

arr = [5, 0, 5, 10, 3, 2, -15, 4]
n = len(arr)
k = 5

# variable to store count of subarrays
res = 0

# Calculate all subarrays
for i in range(n):
	sum = 0  
	for j in range(i, n):
 		# Calculate required sum
 		sum += arr[j]
 		
 		# Check if sum is equal to required sum
 		if sum == k:
  			res += 1

# print the result
print(res)

Output

Complexity Analysis

Time Complexity

O(N^2), where ‘N’ is the number of elements in the array/list ‘arr’.

Since we are iterating over all the elements of the array/list ‘arr’, counting how many subarrays present in ‘ARR’ have sum equal to ‘k’ and increment the ‘res’ accordingly hence, the complexity here grows by O(N^2).

Space Complexity

O(1). Since, Here, we are using constant space therefore space complexity is O(1).

Approach 2: Hashing

An efficient solution is while traversing the array, we keep storing the sum that we have gotten so far in ‘currsum’. Here would also need to maintain the count of different values of ‘currsum’ in a map. Now, if the value of ‘currsum’ equals the desired sum at any instance we increment the count of subarrays by one.

While if the value of ‘CURRSUM’ exceeds the desired sum by a value (‘currsum’ - ‘totalsum’), then if this value is removed from ‘currsum’ , the desired sum can be obtained. From the map find the number of subarrays previously found having a sum equal to (‘currsum’ - ‘totalsum’). Excluding all those subarrays from the current subarray gives new subarrays having the desired sum.

The algorithm to calculate the required will be:-

Initialize a ‘currsum’ variable with 0, ‘res’ by 0, and a map to store all the count of subarrays having a sum equal to a particular ‘currsum’ that we have created.
If we iterate through the ‘arr’ array and keep finding out the value of ‘currsum’ and check if it equals the value of ‘k’, we increment the value of ‘res’.
If the value of ‘currsum’ exceeds the desired sum (which would be by a value (‘currsum’ - ‘totalsum’). So increase count by the number of such subarrays and remove this value from ‘currsum’, which would give us the desired sum.
- In such a case, from the map that we have obtained, find the number of subarrays previously found having sum equal to (‘currsum’ - ‘totalsum’). Excluding all those subarrays from the current subarray, gives new subarrays having the desired sum.
Note that when ‘currsum’ is equal to the desired sum then also we should check the number of subarrays previously having a sum equal to 0. Excluding those subarrays from the current subarray gives new subarrays having the desired sum. Here also we would increase the value of count by the number of subarrays having sum 0.
Finally, we can return the ‘res’

C++ Code

/*
    Time Complexity: O(N)
    Space Complexity: O(N)

    Where 'N' is the number of elements in the array/list 'arr'.
*/

int findAllSubarraysWithGivenSum(vector < int > & arr, int k) {

    // Used to store a number of subarrays starting from index zero having a particular value of sum.
    unordered_map < int, int > prevSum;
    int n = arr.size();
    int res = 0;

    // To keep track of the sum of elements so far.
    int currsum = 0;

    for (int i = 0; i < n; i++) {

        // Add the current element to sum so far.
        currsum += arr[i];

        // If the current sum is equal to the desired sum,
        // then a new subarray is found. 
        // So, increase the count of subarrays.
        if (currsum == k) {
            res += 1;
        }

        /*
            If current sum exceeds given sum by current sum  - sum then
            find the number of subarrays having this sum in our map and exclude these subarrays.
        */
        if (prevSum.find(currsum - k) != prevSum.end()) {
            res += prevSum[currsum - k];
        }

        // Add current sum value to count of different values of sum.
        prevSum[currsum] += 1;
    }

    return res;
}

// main program to implement above functions
int main()
{
  	vector <int> arr = {5,0,5,10,3,2,-15,4};
  	int k = 5;
  	
  	cout<<"Number of subarrays with sum equal to "<<k<<" = "<<findAllSubarraysWithGivenSum(arr,k)<<endl;
  	
 	return 0;
}

Java code

// Java program to find number of subarrays
// with sum exactly equal to k.
import java.io.*;
import java.util.HashMap;
import java.util.Map;

public class Solution {
    public static void main(String[] args)
	{
		int arr[] = {5, 0, 5, 10, 3, 2, -15, 4};
		int k = 5;
		int n = arr.length;

		// function to find the number of subarrays
		// with sum k
		int res=findSubarraySum(arr, n, k);

		// print the result
		System.out.println(res);
	}

	// Function to find the number of 
	// subarrays with sum k
	static int findSubarraySum(int arr[], int n, int k)
	{
		// Used to store a number of subarrays starting 
		// from index zero having a particular value of sum
		HashMap<Integer, Integer> prevSum = new HashMap<>();

		// add count of subarrays with 0 sum to be 1
		prevSum.put(0,1);
		int res = 0;

		// to keep track of elements so far
		int currSum = 0;

		for (int i = 0; i < n; i++) {
			// Add current element to sum so far
			currSum += arr[i];

			// Calculate the sum to be subtracted to
			// get the target sum
			int removeSum=currSum-k;
			
			// Find the count of removeSum which occured before
			// and add it to the result
			if (prevSum.containsKey(removeSum)){
    			res += prevSum.get(removeSum);   
			}

			// add the count of currsum in the HashMap
			prevSum.put(currSum,prevSum.getOrDefault(currSum,0)+1);
		}
		
		// return the result
		return res;
	}
}

Python code

# Python3 program to find the number of
# subarrays with sum k
from collections import defaultdict

# Function to find number of subarrays
# with sum k
def findSubarraySum(arr, n, k):

	# DIctionary used to store a number of subarrays starting
	# from index zero having a particular value of sum
	prevSum = defaultdict(lambda : 0)
    	
    # res stores our result
	res = 0

	# Sum of elements so far.
	currsum = 0

	for i in range(0, n):

		# add current element to sum so far
		currsum += arr[i]

		# if currsum is equal to desired sum,
		# then a new subarray is found. So
		# increase the count of subarrays.
		if currsum == k:
			res += 1
			
		# the sum to be subtracted to get 
		# the target sum is (currsum - sum)
		# Find number of subarrays having
		# this sum and add it to the result
		if (currsum - k) in prevSum:
			res += prevSum[currsum - k]

		# increase currsum count in the
		# prevSum dictionary
		prevSum[currsum] += 1

	# return the result
	return res

if __name__ == "__main__":

	arr = [5, 0, 5, 10, 3, 2, -15, 4]

	# k is the given target sum
	k = 5

	n = len(arr)
	print(findSubarraySum(arr, n, k))

Output

Complexity Analysis

Time Complexity

O(N), where ‘N’ is the number of elements in the array/list ‘arr’.

Since we are iterating over all the elements of the array/list ‘arr’ to calculate the prefix sums and check their sums, the complexity here grows by O(N).

Space Complexity

O(N), where ‘N’ is the number of elements in the array/list ‘arr’.

Here, we require an O(N) amount of space for storing the respective sums in the map for indexes. Therefore, the space complexity is O(N).

Frequently Asked Questions

What are subarrays?

A subarray is defined as a contiguous block of elements in the array.

What is the best time complexity to find all the subarrays with a given sum?

The best time complexity to find all the subarrays with a given sum is O(N), where ‘N’ is the number of elements in the array/list ‘arr’.

How many subarrays are there in an array?

There is a total of N*(N+1)/2 subarrays in an array where ‘N’ is the number of elements in the array/list ‘arr’.

Conclusion

Here we learned one of the most important and frequently asked questions in Amazon, Microsoft, and other product-based companies, i.e. Count Number of Subarrays with Sum K. We discussed various approaches ranging from brute force to optimal along with the C++ code for your better understanding.

Some other interesting array problems frequently encountered in interviews are Sort 0,1,2, Count all sub-arrays having sum divisible by k, Maximum subarray sum after K concatenation, and many more.
Also check out some of the Guided Paths on topics such as Data Structure and Algorithms, Competitive Programming, etc. as well as some Contests and some more Interview Experiences curated by top Industry Experts only on CodeStudio.