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

filter_none

edit
close

play_arrow

link
brightness_4
code

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

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

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

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# 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

chevron_right


C#

filter_none

edit
close

play_arrow

link
brightness_4
code

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

chevron_right


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.