kmjp's blog

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

エクサウィザーズ 2019 : E - Black or White

これ系の問題少しは慣れてきたと思いたい。
https://atcoder.jp/contests/exawizards2019/tasks/exawizards2019_e

問題

B個の黒いチョコとW個の白いチョコがある。
これらを1列に並べる。
その際、以下の手順で端から並べていく。

  • 白黒両方のチョコがともに1個以上ある場合、1/2の確率でどちらかを選択して並べる。
  • 片方の色のチョコしか残っていない場合、それを並べる。

i番目に並ぶチョコが黒である確率P(i)を求め、i=1~(B+W)に対し列挙せよ。

解法

f(x) := x個目に並べたチョコでちょうど黒いチョコがなくなる確率
g(x) := x個目に並べたチョコでちょうど白いチョコがなくなる確率
とする。とすると、x個チョコを並べた時点でまだ両方1個以上残っている確率h(x)は
h(x)=1-sum(f(1)~f(x-1))-sum(g(1)~g(x-1))
となる。
この時、x個目が黒い確率は

  • x個目で最後の黒を引く確率f(x)
  • x-1個目までで白がなくなる確率sum(g(1)~g(x-1))
  • x-1個目までで白黒両方残っている場合の1/2の確率h(x)/2

の総和で

p(x) = f(x) + sum(g(1)~g(x-1)) + h(x)/2
となる。
f(x)は(x-1)個目までにB-1個の黒チョコが出てx個目にB個目の黒チョコが出る確率なので、(x-B)<Wの場合のみ正となり f(x) = (1/2)^{x-1}{}_{x-1} C_{B-1}となる。
g(x)も同様に求められて、あとは累積和を使えばp(x)が求められる。

int B,W,re;
ll mo=1000000007;

ll comb(ll N_, ll C_) {
	const int NUM_=1400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

ll ret[202020],ret2[202020];

ll E[202020][3];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>B>>W;
	
	for(i=0;i<W;i++) {
		ll pat=modpow(modpow(2,(i+B)))*comb(i+B-1,B-1)%mo;
		E[i+B][0]=pat;
	}
	for(i=0;i<B;i++) {
		ll pat=modpow(modpow(2,(i+W)))*comb(i+W-1,W-1)%mo;
		E[i+W][1]=pat;
		ret2[i+W+1]=pat;
		
	}
	
	ll R[2]={};
	ll pat=1;
	for(i=1;i<=B+W;i++) {
		(E[i][2]=1+(mo-R[0])+(mo-R[1]))%=mo;
		(R[0]+=E[i][0])%=mo;
		(R[1]+=E[i][1])%=mo;
		
		(ret2[i]+=ret2[i-1])%=mo;
		(ret[i]+=modpow(2)*E[i][2]+ret2[i])%=mo;
		cout<<ret[i]<<endl;
		
	}
	
	
}

まとめ

幸い似た問題を考えたことがあったのですんなり。