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

• Difficulty Level : Expert
• Last Updated : 25 Aug, 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

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

## Javascript

 ``

Output

`7`

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

My Personal Notes arrow_drop_up