Count of subarrays having sum equal to its length - GeeksforGeeks

18 May, 2021
Given an array arr[] of size N, the task is to find the number of subarrays having the sum of its elements equal to the number of elements in it.
Examples:
Input: N = 3, arr[] = {1, 0, 2}Output: 3Explanation:Total number of subarrays are 6 i.e., {1}, {0}, {2}, {1, 0}, {0, 2}, {1, 0, 2}.Out of 6 only three subarrays have the number of elements equals to sum of its elements i.e., 1) {1}, sum = 1, length = 1.2) {0, 2}, sum = 2, length = 2.3) {1, 0, 2}, sum = 3, length = 3.
Input: N = 3, arr[] = {1, 1, 0}Output: 3Explanation:Total number of subarrays are 6 i.e. {1}, {1}, {0}, {1, 1}, {1, 0}, {1, 1, 0}.Out of 6 only three subarrays have the number of elements equals to sum of its elements i.e., 1) {1}, sum = 1, length = 1.2) {1}, sum = 1, length = 1.3) {1, 1}, sum = 2, length = 2.
Naive Approach: The idea is to generate all the subarrays of the array and if the sum of elements of the subarray is equal to the number of elements in it then count this subarray. Print the count after checking all the subarrays.
Time Complexity: O(N2)Auxiliary Space: O(N)
Efficient Approach: This problem can be converted into a simpler problem by using observation. If all the elements of the array are decremented by 1, then all the subarrays of array arr[] with a sum equal to its number of elements are same as finding the number of subarrays with sum 0 in the new array(formed by decrementing all the elements of arr[ ] by 1). Below are the steps:
Decrement all the array elements by 1.Initialize a prefix array with prefix[0] = arr[0].Traverse the given array arr[] from left to right, starting from index 1 and update a prefix sum array as pref[i] = pref[i-1] + arr[i].Initialize the answer to 0.Iterate the prefix array pref[] from left to right and increment the answer by the value of the current element in the map.Increment the value of the current element in the map.Print the value of answer after the above steps.
Decrement all the array elements by 1.
Initialize a prefix array with prefix[0] = arr[0].
Traverse the given array arr[] from left to right, starting from index 1 and update a prefix sum array as pref[i] = pref[i-1] + arr[i].
Initialize the answer to 0.
Iterate the prefix array pref[] from left to right and increment the answer by the value of the current element in the map.
Increment the value of the current element in the map.
Print the value of answer after the above steps.
Below is the implementation of the above approach:
C++
Java
Python3
C#
Javascript
// 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){    // Decrementing all the elements    // of the array by 1    for (int i = 0; i < N; i++)        arr[i]--;     // Making prefix sum array    int pref[N];    pref[0] = arr[0];     for (int i = 1; i < N; i++)        pref[i] = pref[i - 1] + arr[i];     // Declare map to store count of    // elements upto current element    map<int, int> mp;    int answer = 0;     // To count all the subarrays    // whose prefix sum is 0    mp[0]++;     // Iterate the array    for (int i = 0; i < N; i++) {         // Increment answer by count of        // current element of prefix array        answer += mp[pref[i]];        mp[pref[i]]++;    }     // Return the answer    return answer;} // Driver Codeint main(){    // Given array arr[]    int arr[] = { 1, 1, 0 };    int N = sizeof arr / sizeof arr[0];     // Function call    cout << countOfSubarray(arr, N);    return 0;}
// Java program for the above approachimport java.util.*;class GFG{ // Function that counts the subarrays// with sum of its elements as its lengthstatic int countOfSubarray(int arr[], int N){    // Decrementing all the elements    // of the array by 1    for (int i = 0; i < N; i++)        arr[i]--;     // Making prefix sum array    int []pref = new int[N];    pref[0] = arr[0];     for (int i = 1; i < N; i++)        pref[i] = pref[i - 1] + arr[i];     // Declare map to store count of    // elements upto current element    HashMap<Integer,              Integer> mp = new HashMap<Integer,                                        Integer>();    int answer = 0;     // To count all the subarrays    // whose prefix sum is 0    mp.put(0, 1);     // Iterate the array    for (int i = 0; i < N; i++)    {         // Increment answer by count of        // current element of prefix array         if(mp.containsKey(pref[i]))        {            answer += mp.get(pref[i]);            mp.put(pref[i], mp.get(pref[i]) + 1);                     }        else        {            mp.put(pref[i], 1);        }    }     // Return the answer    return answer;} // Driver Codepublic static void main(String[] args){    // Given array arr[]    int arr[] = { 1, 1, 0 };    int N = arr.length;     // Function call    System.out.print(countOfSubarray(arr, N));}} // This code is contributed by sapnasingh4991
# 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):     # Decrementing all the elements    # of the array by 1    for i in range(N):        arr[i] -= 1     # Making prefix sum array    pref = [0] * N    pref[0] = arr[0]     for i in range(1, N):        pref[i] = pref[i - 1] + arr[i]     # Declare map to store count of    # elements upto current element    mp = defaultdict(lambda : 0)    answer = 0     # To count all the subarrays    # whose prefix sum is 0    mp[0] += 1     # Iterate the array    for i in range(N):         # Increment answer by count of        # current element of prefix array        answer += mp[pref[i]]        mp[pref[i]] += 1     # Return the answer    return answer # Driver Code # Given array arr[]arr = [ 1, 1, 0 ]N = len(arr) # Function callprint(countOfSubarray(arr, N)) # This code is contributed by Shivam Singh
// 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 int countOfSubarray(int []arr, int N){    // Decrementing all the elements    // of the array by 1    for (int i = 0; i < N; i++)        arr[i]--;     // Making prefix sum array    int []pref = new int[N];    pref[0] = arr[0];     for (int i = 1; i < N; i++)        pref[i] = pref[i - 1] + arr[i];     // Declare map to store count of    // elements upto current element    Dictionary<int,               int> mp = new Dictionary<int,                                        int>();    int answer = 0;     // To count all the subarrays    // whose prefix sum is 0    mp.Add(0, 1);     // Iterate the array    for (int i = 0; i < N; i++)    {         // Increment answer by count of        // current element of prefix array         if(mp.ContainsKey(pref[i]))        {            answer += mp[pref[i]];            mp[pref[i]]= mp[pref[i]] + 1;           }        else        {            mp.Add(pref[i], 1);        }    }     // Return the answer    return answer;} // Driver Codepublic static void Main(String[] args){    // Given array []arr    int []arr = { 1, 1, 0 };    int N = arr.Length;     // Function call    Console.Write(countOfSubarray(arr, N));}} // This code is contributed by sapnasingh4991
<script>// Js program for the above approach // Function that counts the subarrays// with sum of its elements as its lengthfunction countOfSubarray( arr, N){    // Decrementing all the elements    // of the array by 1    for (let i = 0; i < N; i++)        arr[i]--;     // Making prefix sum array    let pref = [];    pref[0] = arr[0];     for (let i = 1; i < N; i++)        pref[i] = pref[i - 1] + arr[i];     // Declare map to store count of    // elements upto current element    let mp = new Map;    let answer = 0;     // To count all the subarrays    // whose prefix sum is 0    if(mp[0])    mp[0]++;    else    mp[0] = 1;     // Iterate the array    for (let i = 0; i < N; i++) {        // Increment answer by count of        // current element of prefix array        if(mp[pref[i]]){        answer += mp[pref[i]];        mp[pref[i]]++;        }    }     // Return the answer    return answer;} // Driver Code// Given array arr[]let arr = [ 1, 1, 0 ];let N = arr.length;// Function calldocument.write(countOfSubarray(arr, N)); </script>
3
 
Time Complexity: O(N * Log(N))Auxiliary Space: O(N)
sapnasingh4991
rohitsingh07052
cpp-map
subarray
subarray-sum
Arrays
Competitive Programming
Dynamic Programming
Greedy
Mathematical
Arrays
Dynamic Programming
Greedy
Mathematical
Writing code in comment?
Please use ide.geeksforgeeks.org,
generate link and share the link here.
Maximum and minimum of an array using minimum number of comparisons
Top 50 Array Coding Problems for Interviews
Multidimensional Arrays in Java
Introduction to Arrays
Linear Search
Competitive Programming - A Complete Guide
Practice for cracking any coding interview
Arrow operator -> in C/C++ with Examples
Top 10 Algorithms and Data Structures for Competitive Programming
Bits manipulation (Important tactics)