# Count of subarrays having sum equal to its length | Set 2

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

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 ` `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++;`   `    ``// 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;`   `    ``// Function Call` `    ``countOfSubarray(arr, N);`   `    ``return` `0;` `}`

## Java

 `// 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 mp = ``new` `HashMap();  `   `    ``// 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

 `# 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#

 `// 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;`   `  ``// 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`

Output

`7`

Time Complexity: O(N)
Auxiliary Space: O(N)

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.

My Personal Notes arrow_drop_up Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.