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

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(N²)
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)