yukicoder : No.2457 Stampaholic (Easy)

Hard解いたらEasyも同じコードで通るかと思ったら、微妙に制限が違った。
https://yukicoder.me/problems/no/2457

問題

整数H,W,N,Kが与えられる。
H*Wのグリッドがあり、初期状態で各マスは白色である。
ここに、K*Kのサイズの黒いスタンプを、マス目に沿ってどこかの領域に等確率で押すことをN回繰り返す。
最終的な黒マス数の期待値を求めよ。

解法

あるマスについて、スタンプの押し方のバリエーションがS通りあり、そのうちそのマスを黒くするような塗り方がM通りあれば、そのマスが黒い確率P(M)は
である。

マスを黒くするような塗り方がM通りであるマスの個数とcnt(M)とすると、求めるのは

である。Mは高々K^2通りなので、各Mに対しcnt(M)を求めればよい。

あるマスに対し、そのマスを黒くできるようなスタンプの位置の行の候補がh個ある行をch(h)、列の候補がh個ある列をcw(w)とすると、
である。
ch(h)、cw(w)はO(1)で計算できるので、h,wを1～Kまで総当たりすればcnt(M)を計算できる。

const int 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_=2400001;
	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;
}

int H,W,N,K;


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	
	cin>>H>>W>>N>>K;
	
	vector<ll> Hs(K+1),Ws(K+1);
	
	int numH=0,Hma=0;;
	int numW=0,Wma=0;;
	if(2*K<=H) {
		for(i=1;i<K;i++) Hs[i]=2;
		Hs[K]=H-2*K+2;
	}
	else {
		for(i=1;i<H-K+1;i++) Hs[i]=2;
		Hs[H-K+1]=2*K-H;
	}
	if(2*K<=W) {
		for(i=1;i<K;i++) Ws[i]=2;
		Ws[K]=W-2*K+2;
	}
	else {
		for(i=1;i<W-K+1;i++) Ws[i]=2;
		Ws[W-K+1]=2*K-W;
	}
	
	ll ret=0;
	ll L=1LL*(H-K+1)*(W-K+1)%mo;
	ll rL=modpow(L);
	
	for(x=1;x<=K;x++) for(y=1;y<=K;y++) {
		ret+=Hs[x]*Ws[y]%mo*(1-modpow((L-x*y)*rL,N))%mo;
	}
	
	ret=(ret%mo+mo)%mo;
	cout<<ret<<endl;
	
}

まとめ

色々解法ありそう。