Given an array arr[] of size N, the task is to find the number of subarrays having sum of its elements equal to the number of elements in it.
Examples:
Input: N = 3, arr[] = {1, 0, 2}
Output: 3
Explanation:
Total number of subarrays are 6 i.e., {1}, {0}, {2}, {1, 0}, {0, 2}, {1, 0, 2}.
Out of the 6 subarrays, following three subarrays satisfy the given conditions:
- {1}: Sum = 1, Length = 1
- {0, 2}: Sum = 2, Length = 2
- {1, 0, 2}: Sum = 3, Length = 3
Input: N = 3, arr[] = {1, 1, 0}
Output: 3
Explanation:
Total number of subarrays are 6, i.e. {1}, {1}, {0}, {1, 1}, {1, 0}, {1, 1, 0}.
Out of the 6 subarrays, following three subarrays satisfy the given conditions:
- {1}: Sum = 1, Length = 1
- {1, 1}: Sum = 2, Length = 2
- {1}: Sum = 1, Length = 1
Input: N = 6, arr[] = {1, 1, 1}Output: 3Explanation:Total number of subarrays are 6, i.e. {1}, {1}, {1}, {1, 1}, {1, 1}, {1, 1, 1}.Out of the 6 subarrays, following six subarrays satisfy the given conditions:
- {1},{1},{1}: Sum=1, Length=1 ,Number of Subarrays=3
- {1,1},{1,1}: Sum=2, Length=2, Number of Subarrays=2
- {1,1,1}: Sum=3, Length=3, Number of Subarrays=1
Total Number of subarrays having Sum equals to it’s Length=6
Naive and Prefix Sum based Approach: Refer to previous post for the simplest and prefix sum based approaches to solve the problem.
Time Complexity: O(N2)
Auxiliary Space: O(N)
Efficient Approach: To optimize the above approach, the idea is to store the previous occurrences of subarrays with the given conditions and make use of unordered_map for constant lookup. Below are the steps:
- Initialize an unordered_map M, answer to store the count of subarrays, and sum to store the prefix sum of the array.
-
Traverse the given array and do the following:
- Add the current element to the sum.
- If M[sum – i] exists then add this value to the answer as there exists a subarray of length i whose sum of the element is the current sum.
- Increment the frequency of (sum – i) in the map M.
- After the above steps, print the value of the answer as the total count of subarrays.
Below is the implementation of the above approach:
// C++ program for the above approach #include <bits/stdc++.h> using namespace std;
// Function that counts the subarrays // with sum of its elements as its length int countOfSubarray( int arr[], int N)
{ // Store count of elements upto
// current element with length i
unordered_map< int , int > mp;
// Stores the final count of subarray
int answer = 0;
// Stores the prefix sum
int sum = 0;
// If size of subarray is 1
mp[1]++;
// Iterate the array
for ( int i = 0; i < N; i++) {
// Find the sum
sum += arr[i];
answer += mp[sum - i];
// Update frequency in map
mp[sum - i]++;
}
// Print the total count
cout << answer;
} // Driver Code int main()
{ // Given array arr[]
int arr[] = { 1, 0, 2, 1, 2, -2, 2, 4 };
// Size of array
int N = sizeof arr / sizeof arr[0];
// Function Call
countOfSubarray(arr, N);
return 0;
} |
// Java program for the above approach import java.util.*;
class GFG{
// Function that counts the subarrays // with sum of its elements as its length static void countOfSubarray( int arr[], int N)
{ // Store count of elements upto
// current element with length i
Map<Integer,
Integer> mp = new HashMap<Integer,
Integer>();
// Stores the final count of subarray
int answer = 0 ;
// Stores the prefix sum
int sum = 0 ;
// If size of subarray is 1
if (mp.get( 1 ) != null )
mp.put( 1 , mp.get( 1 ) + 1 );
else
mp.put( 1 , 1 );
// Iterate the array
for ( int i = 0 ; i < N; i++)
{
// Find the sum
sum += arr[i];
if (mp.get(sum - i) != null )
answer += mp.get(sum - i);
// Update frequency in map
if (mp.get(sum - i) != null )
mp.put(sum - i, mp.get(sum - i) + 1 );
else
mp.put(sum - i, 1 );
}
// Print the total count
System.out.print(answer);
} // Driver Code public static void main(String args[])
{ // Given array arr[]
int arr[] = { 1 , 0 , 2 , 1 , 2 , - 2 , 2 , 4 };
// Size of array
int N = arr.length;
// Function Call
countOfSubarray(arr, N);
} } // This code is contributed by ipg2016107 |
# Python3 program for the above approach from collections import defaultdict
# Function that counts the subarrays # with sum of its elements as its length def countOfSubarray(arr, N):
# Store count of elements upto
# current element with length i
mp = defaultdict( lambda : 0 )
# Stores the final count of subarray
answer = 0
# Stores the prefix sum
sum = 0
# If size of subarray is 1
mp[ 1 ] + = 1
# Iterate the array
for i in range (N):
# Find the sum
sum + = arr[i]
answer + = mp[ sum - i]
# Update frequency in map
mp[ sum - i] + = 1
# Print the total count
print (answer)
# Driver code if __name__ = = '__main__' :
# Given array
arr = [ 1 , 0 , 2 , 1 , 2 , - 2 , 2 , 4 ]
# Size of the array
N = len (arr)
# Function Call
countOfSubarray(arr, N)
# This code is contributed by Shivam Singh |
// C# program for the // above approach using System;
using System.Collections.Generic;
class GFG{
// Function that counts // the subarrays with sum // of its elements as its length static void countOfSubarray( int []arr,
int N)
{ // Store count of elements upto
// current element with length i
Dictionary< int ,
int > mp = new Dictionary< int ,
int >();
// Stores the readonly
// count of subarray
int answer = 0;
// Stores the prefix sum
int sum = 0;
// If size of subarray is 1
mp[1] = 1;
// Iterate the array
for ( int i = 0; i < N; i++)
{
// Find the sum
sum += arr[i];
if (mp.ContainsKey(sum - i))
answer += mp[sum - i];
// Update frequency in map
if (mp.ContainsKey(sum - 1))
mp[sum - 1]++;
else
mp[sum - 1] = 1;
}
// Print the total count
Console.Write(answer - 2);
} // Driver Code public static void Main(String []args)
{ // Given array []arr
int []arr = {1, 0, 2, 1,
2, -2, 2, 4};
// Size of array
int N = arr.Length;
// Function Call
countOfSubarray(arr, N);
} } // This code is contributed by gauravrajput1 |
<script> // JavaScript program for the above approach // Function that counts the subarrays // with sum of its elements as its length function countOfSubarray(arr, N)
{ // Store count of elements upto
// current element with length i
var mp = new Map();
// Stores the final count of subarray
var answer = 0;
// Stores the prefix sum
var sum = 0;
// If size of subarray is 1
if (!mp.has(1))
mp.set(1, 1)
else
mp.set(1, mp.get(1)+1)
// Iterate the array
for ( var i = 0; i < N; i++) {
// Find the sum
sum += arr[i];
answer += mp.has(sum - i)?mp.get(sum - i):0;
// Update frequency in map
if (mp.has(sum - i))
mp.set(sum - i, mp.get(sum - i)+1)
else
mp.set(sum - i, 1)
}
// Print the total count
document.write( answer);
} // Driver Code // Given array arr[] var arr = [1, 0, 2, 1, 2, -2, 2, 4];
// Size of array var N = arr.length;
// Function Call countOfSubarray(arr, N); </script> |
7
Time Complexity: O(N)
Auxiliary Space: O(N)