kmjp's blog

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

AtCoder ARC #136 : F - Flip Cells

ARC

こういうの思いつく気がしないなぁ。
https://atcoder.jp/contests/arc136/tasks/arc136_f

問題

H*Wのグリッドがあり各マス0/1が書かれている。
1つのマスを等確率でランダムに選び、0/1反転することを繰り返す。
各行rの総和がA[r]に一致した時点で終了するとする。
終了までに行う反転回数の期待値を求めよ。

解法

考察部分はEditorialに任せて実装だけ。
以下の母関数を考える。
E(x) := 与えられた初期状態から初めて、n次の係数が反転n回で初めて条件を満たす状態に至る確率
F(x) := 与えられた初期状態から初めて、n次の係数が反転n回で(初めてでもどうでも良いので)条件を満たす状態に至る確率
G(x) := 条件を満たした状態から、n次の係数が反転n回でまた条件を満たす状態に戻っている確率

E(x)*G(x)=F(x)とすると、E'(1)が解となるので、E'(x)=(F(x)/G(x))'を求める。
ただ、後の手順でF(x)やG(x)は1/(1-x)が登場するので、あらかじめ(1-x)を掛けておき、E'(x)=(((1-x)F(x))/*1^e*(exp(x/HW)-exp(-x/HW))^o/(2^(HW))
となる。Gexp(x)も同様。
上記総和は、まずe・oに対するnをDPで求めた後、e,oに対応して上記多項式を展開する。

Fexp(x)のexp(ax)の部分を1/(1-ax)に置き換えるとF(x)になる。
これでF(x)やG(x)は(1-x)/(1-ax)の線形和になるので、微分した値も容易に計算可能になる。

int H,W;
int A[55],B[55];
ll Ffrom[2555],Fto[2555],Gfrom[2555],Gto[2555],F[2555],G[2555];
ll P[2525];
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;
}

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	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;
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>H>>W;
	FOR(y,H) {
		cin>>s;
		FORR(c,s) if(c=='1') B[y]++;
	}
	FOR(y,H) {
		cin>>A[y];
	}
	Ffrom[0]=Gfrom[0]=1;
	FOR(y,H) {
		ZERO(Fto);
		ZERO(Gto);
		
		int e1,e0,o1,o0;
		FOR(e1,B[y]+1) {
			o1=B[y]-e1;
			o0=A[y]-e1;
			e0=W-e1-o0-o1;
			if(o1<0||o0<0||e0<0) continue;
			ll pat=comb(e1+o1,e1)*comb(e0+o0,e0)%mo;
			FOR(i,2501) (Fto[i+e1+e0]+=pat*Ffrom[i])%=mo;
		}
		FOR(e1,A[y]+1) {
			o1=o0=A[y]-e1;
			e0=W-e1-o0-o1;
			if(o1<0||o0<0||e0<0) continue;
			ll pat=comb(e1+o1,e1)*comb(e0+o0,e0)%mo;
			FOR(i,2501) (Gto[i+e1+e0]+=pat*Gfrom[i])%=mo;
		}
		
		swap(Ffrom,Fto);
		swap(Gfrom,Gto);
	}
	
	ll p2=1;
	FOR(i,H*W) p2=p2*modpow(2)%mo;
	FOR(i,H*W+1) P[i]=comb(H*W,i)*p2%mo;
	
	for(i=H*W;i>=0;i--) {
		
		FOR(j,H*W+1) {
			(F[j]+=Ffrom[i]*P[j])%=mo;
			(G[j]+=Gfrom[i]*P[j])%=mo;
		}
		
		if(i) {
			//x+1で割る
			FOR(j,H*W) {
				(P[j+1]+=mo-P[j])%=mo;
			}
			assert(P[H*W]==0);
			for(j=H*W-1;j>=0;j--) {
				(P[j+1]+=P[j])%=mo;
				(P[j]=mo-P[j])%=mo;
			}
		}
	}
	ll Fv=F[H*W];
	ll Gv=G[H*W];
	ll dFv=0,dGv=0;
	FOR(i,H*W) {
		ll a=(mo-H*W+2*i)*modpow(H*W)%mo;
		a=modpow(a-1);
		(dFv+=F[i]*a)%=mo;
		(dGv+=G[i]*a)%=mo;
	}
	ll p=((dFv*Gv-Fv*dGv)%mo+mo)%mo;
	ll q=Gv*Gv%mo;
	cout<<(p*modpow(q))%mo<<endl;
	
	
	
}

感想

うーん、母関数周りの知識が足りないせいか、この考察がとても自分でできる気がしない。

*1:1-x)G(1)))'を求めよう。 F(x)にしろG(x)にしろ直接求めるのは大変なので、指数型母関数Fexp(x)・Gexp(x)を考える。 元のH*W個のマスに対し、偶数回処理するものがe個、奇数回処理するものがo個あるときに条件を満たすものがn通りとすると、 Fexp(x) += n*(exp(x/HW)+exp(-x/HW