Skip to content
Related Articles

Related Articles

Improve Article
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 <bits/stdc++.h>
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<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 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




<script>
 
// JavaScript program for the above approach
 
// Function that counts the subarrays
// with sum of its elements as its length
function 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 array
var N = arr.length;
 
// Function Call
countOfSubarray(arr, N);
 
 
</script>
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 Geeks Classes Live 




My Personal Notes arrow_drop_up
Recommended Articles
Page :