kmjp's blog

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

Codeforces ECR #133 : F. Bags with Balls

これ系はだいぶ慣れてきたな。
https://codeforces.com/contest/1716/problem/F

問題

整数N,M,Kが与えられる。

N個のバッグがあり、各バッグには1~M番の番号が書かれたボールが1つずつ、計NM個のボールが入っている。
各バッグから1個ずつボールを選ぶとき、奇数の書かれたボールの数をFとする。
各ボールの取り方に対し、F^Kの総和を求めよ。

解法

各バッグからボールを選んだ時、そのボールから重複込みでK個選んだ時にそれらが全部奇数である組み合わせを求めればよい。
事前に以下を計算しておこう。

f(x,y) := x個の中から順序込み・重複で物を選ぶとき、y個のユニークなものが選ばれる組み合わせ

K個選んだ時にそれらが全部奇数である組み合わせは、i個ユニークな奇数を選ぶ場合、
 \displaystyle P(N,i)floor(\frac{M+1}{2})^iM^{N-i}f(K,i)
の総和を取ればよい。

int T;
ll N,M,K;
const ll mo=998244353;

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

ll from[2020],to[2020];
ll dp[2020][2020];


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	dp[0][0]=1;
	for(i=1;i<=2000;i++) {
		FOR(j,2001) {
			// same
			(dp[i][j]+=dp[i-1][j]*j)%=mo;
			// new
			(dp[i][j+1]+=dp[i-1][j])%=mo;
		}
		swap(from,to);
	}
	
	
	cin>>T;
	while(T--) {
		cin>>N>>M>>K;
		
		
		ll ret=0;
		ll p=1,q=1;
		for(i=1;i<=min(K,N);i++) {
			(p*=N+1-i)%=mo;
			ll b=modpow((M+1)/2,i);
			ll c=modpow(M,N-i);
			ll d=dp[K][i];
			ret+=p*b%mo*c%mo*d%mo;
		}
		cout<<ret%mo<<endl;
	}
}

まとめ

6問回のボス問でこの短さは珍しいね。