C++(g++ 7.5.0) 解法, 执行用时: 145ms, 内存消耗: 19788K, 提交时间: 2023-03-18 17:28:02

#include<bits/stdc++.h>
using namespace std;
int main(){
	int t;
	cin>>t;
	while(t--){
		int n;
		string temp;
		cin>>n>>temp;
		int x=0,y=0;
		map<pair<int,int>,int > cha;
		cha[{0,0}]=1; 
		long long ans=0;
		for(int i=0;i<n;i++){
			if(temp[i]=='0'){
				x++,y++;
			}else if(temp[i]=='1'){
				x--;
			}else y--;
			ans+=cha[{x,y}];
			cha[{x,y}]++;
		}
		cout<<ans<<endl;
	}
	return 0;
}

C++(clang++ 11.0.1) 解法, 执行用时: 102ms, 内存消耗: 19476K, 提交时间: 2023-03-26 12:15:43

#include<bits/stdc++.h>
using namespace std;
int arr[200];
int main()
{
	int t,n;
	string s;
	cin>>t;
	while(t--)
	{
		long long res=0;
		map<pair<int,int>,int> ma;
		memset(arr,0,sizeof(arr));
		cin>>n>>s;
		ma[{0,0}]=1;
		for(int i=0;i<n;i++)
		{
			arr[s[i]]++;
			res+=ma[{arr['2']-arr['1'],arr['1']-arr['0']}]++;
		}
		cout<<res<<endl;
	}
}

Python3 解法, 执行用时: 520ms, 内存消耗: 50548K, 提交时间: 2023-03-18 22:45:45

t=int(input())
for __ in range(t):
    n=int(input())
    s=input()
    cnt=dict(zip("012",[0]*3))
    res={(0,0):1}
    ans=0
    for i in s:
        cnt[i]+=1;
        sta=(cnt['0']-cnt['1'],cnt['1']-cnt['2'])
        ans+=res.get(sta,0)
        if sta in res:res[sta]+=1
        else:res[sta]=1

    print(ans)