Codechef March Cook-Off 2014 - ABC-Strings

問題文 : ABC-Strings

概要 : 文字列 $S$ の部分文字列 $t$ であって, $t$ の中の A の個数, B の個数, C の個数が等しい物の個数を数えよ.

制約 :
$1 \leqq |S| \leqq 10^6$
$S$ は A, B, C だけからなる文字列である.

解法 :
$D[i]$ を, $i$ 番目の文字までの A の個数 - B の個数, B の個数 - C の個数の組とする.

すると, $D[i] = D[j]$ となる点 $i,\ j$ で
$numA[i] = num B[i] + x$, $num B[i] = numC[i] + y$, $numA[j] = num B[j] + x$, $num B[j] = numC[j] + y$が成り立つ.

同じ文字に付いての等式を辺々引いて $numA[i] - numA[j] = num B[i] - num B[j] = numC[i] - numC[j]$ が成り立つ. すなわち $(i, j]$ 区間で numA = numB = numC となる.

ゆえに, $(x, y)$ の組が等しくなる物でグループを作って, それぞれのグループから2 つ取る組み合わせを足しあわせたものが答えになる.

(JOI の春合宿で出題された JOIOJI という問題とほとんど同じ問題だったため, コードを流用しています.)

コード :

#include <bits/stdc++.h>

#define MAX_N (1000001)

using namespace std;

struct Diff {
	int jo, oi, pos;
	Diff(int a, int b, int c) : jo(a), oi(b), pos(c){}
	Diff(){}
	bool operator < (const Diff &x) const {
		if (x.jo != jo) return (jo < x.jo);
		if (x.oi != oi) return (oi < x.oi);
		return (pos < x.pos);
	}
};

int main()
{
	int N;
	static char s[MAX_N];
	static Diff diffs[MAX_N + 1], it = Diff(0, 0, -1);
	
	int idx = 0;
	diffs[idx++] = it;
	
	scanf("%s", s);
	N = strlen(s);
	for (int i = 0; s[i]; i++){
		it.pos++;
		if (s[i] == 'A') it.jo++;
		if (s[i] == 'B') it.jo--, it.oi++;
		if (s[i] == 'C') it.oi--;
		diffs[idx++] = it;
	}
	
	sort(diffs, diffs + idx);
	
	it = diffs[0];
	long long int ret = 0, cnt = 1;
	for (int i = 1; i <= N; i++){
		if (it.jo == diffs[i].jo && it.oi == diffs[i].oi)
			cnt++;
		else {
			it = diffs[i];
			ret += cnt * (cnt - 1) / 2;
			cnt = 1;
		}
	}
	ret += cnt * (cnt - 1) / 2;
	
	printf("%lld\n", ret);
	return (0);
}