ZR1001 欢乐 ABC（map 维护前缀和）

Description

求字符串 $S$ 中有多少个子串，满足 $A$、$B$、$C$ 出现次数相等。

$1\le |S|\le 10^6$

Solution

若有 $A$、$B$、$C$ 出现次数的前缀和。对于一个区间 $[l,r]$，若满足下述条件，则合法。

$$A_r-A_{l-1}=B_r-B_{l-1}⟹A_r-B_r=A_{l-1}-B_{l-1}$$

$$B_r-B_{l-1}=C_r-C_{l-1}⟹B_r-C_r=B_{l-1}-C_{l-1}$$

用 map 维护 $(A_i-B_i,B_i-C_i)$ 的二元组即可，注意 $(0,0)$ 应赋初值为 $1$。

Code

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 1e6 + 5;
int read() {
	int x = 0, f = 0; char ch = 0;
	while (!isdigit(ch)) f |= ch == '-', ch = getchar();
	while (isdigit(ch)) x = (x << 3) + (x << 1) + (ch ^ 48), ch = getchar();
	return f ? -x : x;
}
char s[N];
map < pair<int, int>, int> mp; 
int main(){
	scanf("%s", s);
	int n = strlen(s);
	int A = 0, B = 0, C = 0;
	ll ans = 0; 
	mp[make_pair(0, 0)] = 1;
	for (int i = 0; i < n; i++) {
		if (s[i] == 'A') A++;
		else if (s[i] == 'B') B++;
		else if (s[i] == 'C') C++;
		ans += mp[make_pair(A - B, B - C)]++;
	}
	printf("%lld\n", ans);
	return 0;
}