Codeforces #213 Div1. A. Matrix

今回Aをさっくり解けたと思ったら、まさかのint型オーバーフローで落とした…。
http://codeforces.com/contest/364/problem/A

問題

数値で構成された文字列Sと、整数Aが与えられる。
ここで、B[i][j]=S[i]*S[j]とする。

x<=i<=y,z<=j<=tとなるようなi,jにおいてB[i][j]の和がAとなるような(x,y,z,t)の組み合わせの数を答えよ。

解法

B[i][j]の和は(S[x]+S[x+1]+..+S[y])*(S[z]+S[z+1]+..+S[t])となる。

まず各x,yに対しS[x]+S[x+1]+..+S[y]を求め、その和を構成できる(x,y)の組みをカウントする。
Sの長さをN文字とすると、N<=4000なのでO(N^2)で求めれば間に合う。

各(x,y)における(S[x]+..+S[y])=Pとすると、先のカウントからA/Pとなるような(z,t)の組み合わせの数がわかるので、それを加算していけばよい。

A==0の時は(x,y)と(z,t)のどちらかが0であればよい点に注意。
多くの人がA==0に引っかかっていた。

int A;
string S;
int SS[4005];
map<int,ll> MM;
void solve() {
	int f,i,j,k,l,x,y;
	ll ret;
	
	cin>>A>>S;
	FOR(i,S.size()) SS[i+1]=SS[i]+S[i]-'0';
	for(x=0;x<S.size();x++) for(y=x+1;y<=S.size();y++) MM[SS[y]-SS[x]]++;
	
	ret=0;
	if(A==0) {
		ll n0=0,n1=0;
		ITR(it,MM) {
			if(it->first==0) n0+=it->second;
			else n1+=it->second;
		}
		ret=(n0+n1)*(n0+n1)-n1*n1;
	}
	else {
		for(x=1;x*x<=A;x++) {
			if(A%x) continue;
			if(x*x==A) ret += MM[x]*MM[A/x];
			else ret += 2*MM[x]*MM[A/x];
		}
	}
	cout << ret << endl;
}

まとめ

もうオーバーフローは嫌だ…。