yukicoder : No.1289 RNG and OR

これ系頭混乱するよね。
https://yukicoder.me/problems/no/1289

問題

0～(2^n)-1のいずれかの数値iが手に入るくじがあり、それぞれの確率A[i]が与えられる。
今まで得た数値のorが(2^n)-1になるまでくじを引くとき、くじを引く回数の期待値を求めよ。

解法

E[i]をiに含まれるbitが1個でも手に入っている状態までにくじを引く回数の期待値とする。
一件(E[1]+E[2]+....+E[2^k])が解に思えるが、そこから2bit以上立っているケースを引いて、3bit以上立っているケースを足して…と包除原理の要領で考えていくと、解は

となる。あとはE[i]を求めよう。

iの各ビットを含まない値が出る確率をP[i]とすると、E[i]＝1/(1-P[i])である。
P[i]は高速ゼータ変換で求められる。

int N;
int A[1<<18];
ll SA[1<<18],E[1<<18];
ll S;
const ll mo=998244353;

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;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,1<<N) {
		cin>>A[i];
		S+=A[i];
	}
	FOR(i,1<<N) {
		SA[i]=A[i]*modpow(S)%mo;
	}
	FOR(i,N) {
		FOR(j,1<<N) if(j&(1<<i)) (SA[j]+=SA[j^(1<<i)])%=mo;
	}
	ll ret=0;
	for(i=1;i<1<<N;i++) {
		E[i]=modpow(1+mo-SA[((1<<N)-1)^i]);
		if(__builtin_popcount(i)%2==0) {
			ret-=E[i];
		}
		else {
			ret+=E[i];
		}
	}
	
	cout<<(ret%mo+mo)%mo<<endl;
	
}

まとめ

符号とか包除とか混乱するんだよね。