Codeforces #266 Div2 C. Number of Ways

CF266に参加。Bを落としたけどCDEが何とか解けた。
とはいえミスが多くて微妙なスコア。
http://codeforces.com/contest/466/problem/C

問題

数列A[i]が与えられる。
この数列を連続した部分列3つに分割したとき、それぞれの総和が等しくなるようにしたい。
そのような分割の仕方は何通りあるか。

解法

数列の累積和を求め、まずは総和Sが3の倍数であることを確認する。
後は数列の先頭からの和がS/3の位置を求めて、そのあとにS*2/3となる要素の数を求めていけばよい。

int N;
ll A[600000],S[600000];
int ok[600000];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>A[i];
		S[i+1]=S[i]+A[i];
	}
	
	if(S[N]%3) return _P("0\n");
	for(i=N-2;i>=0;i--) ok[i]=ok[i+1]+(S[i+1]==S[N]*2/3);
	ll ret=0;
	for(i=0;i<N;i++) if(S[i+1]==S[N]/3) ret+=ok[i+1];
	cout << ret << endl;
	
}

まとめ

今回CよりBの方が難しいね。