Description

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

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

Solution

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

$A_r - A_{l-1} = B_r - B_{l-1} \Longrightarrow A_r - B_r = A_{l-1}- B_{l-1}$

$B_r - B_{l-1} = C_r - C_{l-1} \Longrightarrow 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;
}