Educational Codeforces Round 93 (Rated for Div. 2) C. Good Subarrays（前缀和优化技巧）

Posted 2021-05-15 li_wen_zhuo

题目描述

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

Note

The first test case is considered in the statement.
In the second test case, there are 6 good subarrays: a1…1, a2…2, a1…2, a4…4, a5…5 and a4…5.
In the third test case there is only one good subarray: a2…6.

题目大意

给出一个序列（每个数都为0-9），定义好的子序列为：该子序列的和=该子序列的长度。问好的子序列的个数。

题目分析

好的子序列即为：该序列的平均值为1。这样的子序列的和等于该子序列的长度，它的值并不固定，因此我们并不能很好的进行求解。

这里有一个常用的技巧：我们可以给序列中所有数减1。这样我们只需要找出所有序列和为0的子序列来即可，目标这样序列的和就固定了。

除此之外，还有一个技巧：我们可以用map记录下所有前面的前缀和值，如果当前的前缀和值在map中存在，那么说明序列中至少有一段子序列的和值为0.

利用前缀和+map，我们就可以用O(nlogn)的时间内求出所有和为0的子序列了。

代码如下

#include <iostream>
#include <cmath>
#include <cstdio>
#include <set>
#include <string>
#include <cstring>
#include <map>
#include <algorithm>
#include <stack>
#include <queue>
#include <bitset>
#define LL long long
#define ULL unsigned long long
#define PII pair<int,int>
#define x first
#define y second
using namespace std;
const int N=1e3+5,mod=1e9+7;
int main()
{
	int t;
	cin>>t;
	while(t--)
	{
		map<int,LL> mp;
		int n;
		string s;
		cin>>n>>s;
		LL ans=0,sum=0;				//ans记录答案的个数，sum记录前缀和
		for(int i=0;i<s.size();i++)
		{
			int x=s[i]-'1';			//获得当前的数-1
			sum+=x;					//记录前缀和
			if(!sum) ans++;			//如果当前前缀和值=0，说明[0-i]是一个好的子序列
			ans+=mp[sum];			//记录有多少起点不为0的好子序列
			mp[sum]++;				//将当前值放入map
		}
		printf("%lld\\n",ans);
	}
	return 0;
}