总和等于其长度的子数组的计数|套装2

给定大小为N的数组arr [] ，任务是查找其元素之和等于元素数的子数组的数量。

例子：

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

基于朴素和前缀和的方法：有关解决问题的最简单和基于前缀和的方法，请参阅上一篇文章。

时间复杂度： O(N ² )
辅助空间： O(N)

高效方法：为了优化上述方法，其思想是在给定条件下存储先前出现的子数组，并利用unordered_map进行恒定查找。步骤如下：

• 初始化一个unordered_map M ，回答以存储子数组的计数，并求和以存储数组的前缀和。
• 遍历给定数组并执行以下操作：
  • 将当前元素加到sum上。
  • 如果存在M [sum – i] ，则将该值添加到答案中，因为存在一个长度为i的子数组，该数组的元素之和为当前之和。
  • 增大映射M中(sum – i)的频率。
• 完成上述步骤后，将答案的值打印为子数组的总数。

下面是上述方法的实现：

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 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 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 mp = new Dictionary(); 
 
  // 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

输出

7

时间复杂度： O(N)
辅助空间： O(N)