第5回 ドワンゴからの挑戦状 予選 : C - k-DMC

うーん、EはともかくDは解ききりたかった。
https://beta.atcoder.jp/contests/dwacon5th-prelims/tasks/dwacon5th_prelims_c

問題

　
N文字の文字列Sと、クエリKが与えられる。
以下を満たす変数の組は何通りか。

• 0≦a＜b＜c＜N
• c-a＜K
• S[a]='D', S[b]='M', S[c]='C'

解法

以下の2値の累積和を取っておこう。
A(n) := S[0..n]までにおける'C'の登場回数
B(n) := S[0..n]までにおけるS[i]='C'の場合に対するその手前の'M'の登場回数の総和
C(n) := S[0..n]までにおける'M'の登場回数

あとは各クエリに対し、以下の和を取る。

• S[i]='D'の場合、(B(i+K-1)-B(i))-C(i)*(A(i+K-1)-A(i))

要は'D'が登場後そこから右K文字以内に登場する'C'に対し、その手前にある'M'の登場回数の累計を取る。
ただし'D'より手前にある'M'は無視しなければならないので、その分を引いている。

int N;
string S;
int Q,K;
int SM[2010101];
ll BTC[2010101];
ll BTC2[2010101];


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	cin>>S;
	FOR(i,N) {
		SM[i+1]=SM[i]+(S[i]=='M');
		BTC[i+1]=BTC[i];
		BTC2[i+1]=BTC2[i];
		if(S[i]=='C') {
			BTC[i+1]++;
			BTC2[i+1]+=SM[i+1];
		}
	}
	for(i=N+1;i<=2000000;i++) {
		SM[i]=SM[i-1];
		BTC[i]=BTC[i-1];
		BTC2[i]=BTC2[i-1];
	}
	
	cin>>Q;
	while(Q--) {
		cin>>K;
		ll ret=0;
		FOR(i,N) if(S[i]=='D') {
			ll NC=BTC[i+K]-BTC[i];
			ll NC2=BTC2[i+K]-BTC2[i];
			ret+=NC2-NC*SM[i+1];
		}
		cout<<ret<<endl;
	}
	
}

まとめ

BITに行こうとしてしまったのは反省。