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++
// 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 lengthint 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 Codeint 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
// Java program for the above approachimport java.util.*;class GFG{ // Function that counts the subarrays// with sum of its elements as its lengthstatic 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 Codepublic 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
# Python3 program for the above approachfrom collections import defaultdict# Function that counts the subarrays# with sum of its elements as its lengthdef 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 codeif __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#
// C# program for the // above approachusing System;using System.Collections.Generic;class GFG{ // Function that counts // the subarrays with sum // of its elements as its lengthstatic 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 Codepublic 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 |
Javascript
<script>// JavaScript program for the above approach// Function that counts the subarrays// with sum of its elements as its lengthfunction 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 arrayvar N = arr.length;// Function CallcountOfSubarray(arr, N);</script> |
Output
7
Time Complexity: O(N)
Auxiliary Space: O(N)
Feeling lost in the world of random DSA topics, wasting time without progress? It’s time for a change! Join our DSA course, where we’ll guide you on an exciting journey to master DSA efficiently and on schedule.
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
