CSAcademy Round #24 : E. Ball Sampling

これ系苦手意識あったけど解けてよかった。
https://csacademy.com/contest/round-24/#task/ball-sampling

問題

N種類のボールが大量に入った箱がある。
そこから1個取り出したとき、各ボールiの出る確率の比率A[i]が与えられる。
全種類のボールを最低1個ずつ取り出すまでにボールを取り出す回数の期待値を求めよ。

解法

BitDPで解く。
取り出し済みのボールの状態に対して、その状態に至る確率と、その状態に至るボールの回数の期待値を求めよう。
そこからの状態遷移に関して、その状態に至る確率によって期待値を重みづけ平均を取っていけばよい。

int N,A[30];
double prob[1<<20];
double num[1<<20];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	prob[0]=1;
	
	for(int mask=0;mask<1<<N;mask++) {
		int hit=0,miss=0;
		num[mask] /= prob[mask];
		
		FOR(i,N) {
			if(mask&(1<<i)) hit+=A[i];
			else miss+=A[i];
		}
		FOR(i,N) if((mask&(1<<i))==0) {
			prob[mask | (1<<i)] += prob[mask] * A[i]/miss;
			num[mask | (1<<i)] += prob[mask] * A[i]/miss * (num[mask] + 1.0*(hit+miss)/miss);
		}
		
	}
	_P("%.12lf\n",num[(1<<N)-1]);
	
}

まとめ

とはいえ立式には結構手こずった。