Related Articles
Count of subarrays having sum equal to its length | Set 2
• Last Updated : 14 Jun, 2021

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[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

 `// 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] = 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`

## Javascript

 ``
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.  To complete your preparation from learning a language to DS Algo and many more,  please refer Complete Interview Preparation Course.

In case you wish to attend live classes with industry experts, please refer DSA Live Classes

My Personal Notes arrow_drop_up