kmjp's blog

競技プログラミング参加記です

Codeforces #692 Div1 B. Grime Zoo

ちょっとコード量が多いな。
https://codeforces.com/contest/1464/problem/B

問題

0,1,?で構成された文字列Sが与えられる。
S中に部分文字列として"01"があるとコストX、"10"があるとコストYがかかる。
?を0または1に任意に置き換えられるとき、総コストの最小値を求めよ。

解法

?の部分だけ抜き出すと、1111...111000...000か000...000111...111のいずれかのように、0及び1が入る箇所が連続するとよい。
そこで0と1の境目を総当たりしよう。

string S;
ll N,X,Y;
ll L[101010][2],R[101010][2];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>S>>X>>Y;
	N=S.size();
	ll ma=1LL<<60;
	ll base=0;
	FOR(i,N) {
		L[i+1][0]=L[i][0];
		L[i+1][1]=L[i][1];
		if(S[i]=='0') {
			base+=1LL*L[i+1][1]*Y;
			L[i+1][0]++;
		}
		else if(S[i]=='1') {
			base+=1LL*L[i+1][0]*X;
			L[i+1][1]++;
		}
	}
	for(i=N-1;i>=0;i--) {
		R[i+1][0]=R[i+2][0];
		R[i+1][1]=R[i+2][1];
		if(S[i]=='0'||S[i]=='1') R[i+1][S[i]-'0']++;
	}
	ll add=1LL<<60;
	ll sum[2]={};
	int num=0;
	FOR(i,N) if(S[i]=='?') {
		sum[0]+=L[i+1][1]*Y;
		sum[0]+=R[i+1][1]*X;
		sum[1]+=L[i+1][0]*X;
		sum[1]+=R[i+1][0]*Y;
		num++;
	}
	add=min(sum[0],sum[1]);
	int cnt=0;
	FOR(i,N) if(S[i]=='?') {
		//0→1
		sum[0]-=L[i+1][1]*Y;
		sum[0]+=L[i+1][0]*X;
		sum[0]-=R[i+1][1]*X;
		sum[0]+=R[i+1][0]*Y;
		sum[0]+=(num-(cnt+1))*Y;
		sum[0]-=cnt*Y;
		sum[1]-=L[i+1][0]*X;
		sum[1]+=L[i+1][1]*Y;
		sum[1]-=R[i+1][0]*Y;
		sum[1]+=R[i+1][1]*X;
		sum[1]+=(num-(cnt+1))*X;
		sum[1]-=cnt*X;
		add=min(add,sum[0]);
		add=min(add,sum[1]);
		cnt++;
	}
	cout<<base+add<<endl;
}

まとめ

もっとすっきり書けないかな。