Good Subarrays【前缀和】

题目描述：原题链接

You are given an array a1,a2,…,an consisting of integers from 0 to 9. A subarray al,al+1,al+2,…,ar−1,ar is good if the sum of elements of this subarray is equal to the length of this subarray (∑i=lrai=r−l+1).

For example, if a=[1,2,0], then there are 3 good subarrays: a1…1=[1],a2…3=[2,0] and a1…3=[1,2,0].

Calculate the number of good subarrays of the array a.

Input
The first line contains one integer t (1≤t≤1000) — the number of test cases.

The first line of each test case contains one integer n (1≤n≤105) — the length of the array a.

The second line of each test case contains a string consisting of n decimal digits, where the i-th digit is equal to the value of ai.

It is guaranteed that the sum of n over all test cases does not exceed 105.

Output
For each test case print one integer — the number of good subarrays of the array a.

Example
input
3
3
120
5
11011
6
600005
output
3
6
1

思路：

首先求区间和想到前缀和，但是直接求复杂度O(n^2)会超时。我们可以发现，如果 b[i₂] - i₂ = b[i₁] - i₁ —> b[i₂] - b[i₁] = i₂ - i₁ ，也就是如果有两处前缀和减去前缀长度相等，说明它们之间的序列符合要求，有多个的话，每次取一段有(n-1)种，每次取两段有(n-2)种…取n-1段有1种，所以是等差数列求和。注意要把长度直接等于前缀和的加到答案中。

#include <iostream>
#include <algorithm>
#include <cstring>
#include <cmath>
#include <cstdio>
#include <map>
#include <vector>
using namespace std;
typedef long long ll;
const int N=100010;
 
int b[N];
map<int,int> q;
int main()
{
    int T; scanf("%d",&T);
    while(T--)
    {
        int n; scanf("%d",&n);
        char s[N]; scanf("%s",s);
        memset(b,0,sizeof b);
        ll res=0,len=strlen(s);
        for(int i=1;i<=len;i++)
        {
            b[i]=b[i-1]+s[i-1]-'0';
            if(b[i]==i) res++;
        }
        q.clear();
        vector<int> k;
        for(int i=1;i<=n;i++)
        {
            q[b[i]-i]++;
            if(q[b[i]-i]==1) k.push_back(b[i]-i);
        }
        for(int x:k)
        {
            ll tmp=q[x];
            res+=tmp*(tmp-1)/2;
        }
        printf("%lld\n",res);
    }
    return 0;
}