Codeforces #659 Div1 E. Strange Operation

CやDより正答者多い…。
http://codeforces.com/contest/1383/problem/E

問題

01で構成された文字列Sが与えられる。
以下の処理を任意回数行ったとする。
添え字iを指定し、S[i]=max(S[i],S[i+1])で更新したうえで、S[i+1]を削除する。

Sは何通り考えられるか。

解法

Sが0だけで構成されるなら、解は|S|通り。
以下、Sに1が1個以上あるケースを考える。

dp(i) := Sのprefix i文字までの取り込み方を決めたときの、以降の決め方。ただしi文字目のうち'1'が1個以上含まれる

ケースを考える。
dp(i)は、以下の和となる。

• 後ろに1がないなら、ここで文字列を終了する1通り。
• 後ろに1があるなら、それを取り込むケースを考えると、最寄のS[j]=1に対しdp(i)+=dp[j]
• 後ろに0があるなら、それを取り込むケースを考えると、最寄のS[j]=0に対しdp(i)+=dp[j]

解は、S[i]=1となる最小のiに対し、dp(i)*i (最初の1の前の0の個数はi個あるので)となる。

int N;
string S;
ll dp[1010101];
int nex[1010101];
int num0[1010101];
const ll mo=1000000007;

void solve() {
	int i,j,k,l,r,x,y;
	
	cin>>S;
	
	N=S.size();
	if(count(ALL(S),'1')==0) {
		cout<<N<<endl;
		return;
	}
	
	FOR(i,N) {
		if(S[i]=='0') {
			num0[i]=1;
			if(i) num0[i]+=num0[i-1];
		}
	}
	
	FOR(i,N+1) nex[i]=N;
	
	for(i=N-1;i>=0;i--) {
		// next 1
		if(nex[0]<N) dp[i]+=dp[nex[0]];
		// next 0
		if(nex[num0[i]+1]<N) dp[i]+=dp[nex[num0[i]+1]];
		// end
		if(num0[i]<=num0[N-1]) dp[i]++;
		
		dp[i]%=mo;
		nex[num0[i]]=i;
	}
	

	ll ret=0;
	FOR(i,N) if(S[i]=='1') {
		ret=dp[i]*(i+1)%mo;
		break;
	}
	cout<<ret<<endl;
	
}

まとめ

実装は軽いけど、状態遷移で混乱しそう。