CodeChef CookOff 2014 March 3-Palindromes

問題

長さnの数字のみからなる文字列が与えられる。
この数字の連続する部分文字列（ただしleading zeroがあってはならない）で、
回文かつ3の倍数であるものの個数を求めよ。

制約条件

n≦1000000

方針

s[i] = (Σa[i]) % 3とする
i桁目まで、
s[j] = 0となるjがいくつあるか
s[j] = 1となるjがいくつあるか
s[j] = 2となるjがいくつあるか
を累積和を使って求めておく。

Manacherのアルゴリズムで回文半径を求める。

回文半径rが求まったら、中心をiとすれば、[i - k, i + k]（kはr以下）は全て回文なので、
この範囲で3の倍数になるkの個数を、累積和を見て数えればよい。

ソースコード

const int N = 1000100;
int n, rad[2*N], sum[N][3], s[N];
char in[N];

//spaghetti sourceのManacher
void build(char *text, int n) {
  int i, j, k;
  for (i = 0, j = 0; i < 2*n; i += k, j = max(j-k, 0)) {
    while (i-j >= 0 && i+j+1 < 2*n && text[(i-j)/2] == text[(i+j+1)/2]) ++j;
    rad[i] = j;
    for (k = 1; i-k >= 0 && rad[i]-k >= 0 && rad[i-k] != rad[i]-k; ++k)
      rad[i+k] = min(rad[i-k], rad[i]-k);
  }
}

int main(){
	
	gets(in);
	n = strlen(in);
	
	rep(i, n){
		s[i + 1] = s[i] + in[i] - '0';
		s[i + 1] %= 3;
		
		rep(j, 3) sum[i + 1][j] = sum[i][j];
		if(in[i] > '0') sum[i + 1][s[i + 1]]++;
	}
	build(in, n);
	
	ll ans = 0;
	rep(i, 2 * n){
		int l = (i + 1) / 2;
		int r = l + (rad[i] + 1) / 2;
		int mod = i % 2 ? 0 : (in[i / 2] - '0') % 3;
		
		if(mod) mod = 3 - mod;
		mod += s[l];
		mod %= 3;
		
		ans += sum[r][mod] - sum[l][mod];
	}
	rep(i, n) if(in[i] == '0') ans++;
	
	cout << ans << endl;
	
	return 0;
}