补题:牛客NC248906数据结构

传送门:牛客

题目描述:

你有一个长度为n的字符串，其中仅含'0','1','2'三个字符。
你希望知道，这个字符串有多少个子串，满足该子串的'0','1','2'个数相等？
输入:
4
3
012
6
001122
18
102100120120120012
18
012012012012012012
输出:
1
1
25
51

一道小清新思维题.

考虑从左往右进行枚举,遇到0说明我们此时需要一个1和一个2,遇到1说明此时我们需要1的数量减少,遇到2说明此时我们需要2的个数减少.可以使用map来维护1,2需要个数的状态
那么对于前后两个位置(假设为$l$,$r$)假如我们需要1和2的状态是一样的,我们此时的贡献就是前面的状态的贡献+1.因为对于这两个端点中的区间来说,我们从之前的状态经过一系列数字最终回到了这个状态,说明我们的这个区间里面有多个012这样的组合,因为只有这样完整的组合我们的状态才是不会变化的,所以此时我们的这个区间必然是合法的.也就是$[l+1,r]$是一个合法区间(012个数相同).假设我们的$l$及之前有$x$个相同的状态,那么就说明$l$之前有$(x-1)\ast x/2$个贡献,那么对于现在的$r$来说,此时之前的所有相同状态连续到我们现在的$r$都将是一个合法状态,那么此时我们的总贡献就是$x\ast (x+1)/2$,也就是比之前多了一个$x+1$

下面是具体的代码部分:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define root 1,n,1
#define ls rt<<1
#define rs rt<<1|1
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
inline ll read() {
	ll x=0,w=1;char ch=getchar();
	for(;ch>'9'||ch<'0';ch=getchar()) if(ch=='-') w=-1;
	for(;ch>='0'&&ch<='9';ch=getchar()) x=x*10+ch-'0';
	return x*w;
}
#define int long long
#define maxn 1000000
const double eps=1e-8;
#define	int_INF 0x3f3f3f3f
#define ll_INF 0x3f3f3f3f3f3f3f3f
int a[maxn];
map<int,map<int,int> >mp;
signed main() {
	int T=read();
	while(T-- ){
		mp.clear();
		int n=read();
		string s;
		cin>>s;
		for(int i=0;i<s.length();i++) {
			a[i+1]=s[i]-'0';
		}
		int ans=0;
		mp[0][0]=1;
		int x=0,y=0;
		for(int i=1;i<=n;i++) {
			if(a[i]==0) {
				x++;y++;
			}
			else if(a[i]==1) {
				x--;
			}
			else {
				y--;
			}
			ans+=mp[x][y];
			mp[x][y]++;
		}
		cout<<ans<<endl;
	}
	return 0;
}