kmjp's blog

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

Codeforces #854 : G. Count Voting

なんでこれこんな位置にあるんだろ。
https://codeforces.com/contest/1799/problem/G

問題

N人の人がチームに分かれている。
各人が1人に投票するが、その際同チームの人には投票できない。

N人のチーム編成と、各人の得票数が与えられたとき、その状態を満たす投票のパターンは何通りか。

解法

各人がどのチームに投票したかと、チーム内の誰に投票したかは独立に計算して後で掛け合わせればよい。
前者は、
dp(n,k) := nチーム目までの人を考えたとき、k人が(n+1)チーム以降に投票しているような組み合わせの数
としてチーム毎にDPしていくと良い。

int N,C[202],T[202];
vector<int> B[202];
int S[202];
const ll mo=998244353;
ll dp[202][202];
const int NUM_=400001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];

ll comb(ll N_, ll C_) {
	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;
	
	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;
	
	cin>>N;
	FOR(i,N) cin>>C[i];
	FOR(i,N) {
		cin>>T[i];
		B[T[i]-1].push_back(C[i]);
		S[T[i]-1]+=C[i];
	}
	dp[0][0]=1;
	FOR(i,N) {
		FOR(y,201) if(dp[i][y]) {
			FOR(x,min((int)B[i].size(),S[i])+1) {
				(dp[i+1][x+y]+=dp[i][y]*comb(B[i].size(),x)%mo*comb(S[i],x)%mo*fact[x]%mo)%=mo;
			}
		}
	}
	ll ret=0;
	FOR(i,N+1) {
		ll a=dp[N][i]*fact[N-i]%mo;
		if(i%2) {
			ret+=mo-a;
		}
		else {
			ret+=a;
		}
	}
	ret%=mo;
	FOR(i,N+1) if(S[i]) {
		FORR(c,B[i]) ret=ret*factr[c]%mo;
	}
	cout<<ret<<endl;
}

まとめ

実際本番もFより先に解いてるな。