Mujin Programming Challenge 2018 : F - チーム分け

シンプルな設定でいいね。
https://mujin-pc-2018.contest.atcoder.jp/tasks/mujin_pc_2018_f

問題

N人の人をいくつかのグループに分ける。
なお、i番の人はA[i]人以下のグループに属するようにしたい。

グループの分け方は何通りか。

解法

1人ずつグループに追加していこうとすると、各グループの人数を覚えておかなければならず、とてもやっていられない。
そこで、グループを一気に作ることを考える。

以下をiの大きい順に埋めていくことを考えよう。求めたいのはdp(1,0)である。
dp(a,b) := a人以上のグループをいくつか作ったところ、a人以上を許容する人のうちb人がグループ未結成であるような組み合わせ

dp(a+1,b)からの遷移を考える。
A[i]=aとなるiの数をcnt(a)とする。
i人以上を許容する人は(b+cnt(a))人残っていることになるので、ここからa人のグループをc個作るとすると

のように遷移させることができる。

int N;
int A[1010];
ll mo=998244353;
ll dp[1010][1010];

const int NUM_=400001;
ll inv[NUM_+1];
ll combi(ll N, ll C_) {
	static const int N_=1020;
	static ll C[N_][N_];
	int i,j;
	if (fact[0]==0) {
		inv[1]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		
		FOR(i,N_) C[i][0]=C[i][i]=1;
		for(i=1;i<N_;i++) for(j=1;j<i;j++) C[i][j]=(C[i-1][j-1]+C[i-1][j])%mo;
	}
	if(C_<0 || C_>N) return 0;
	return C[N][C_];
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	combi(1,1);
	
	cin>>N;
	FOR(i,N) cin>>x, A[x]++;
	
	dp[1001][0]=1;
	for(i=1000;i>=1;i--) {
		for(j=0;j<=1000;j++) if(dp[i+1][j]) {
			x=j+A[i];
			
			ll pat=dp[i+1][j];
			int left=x;
			for(y=0;y*i<=x;y++) {
				(dp[i][x-y*i]+=pat)%=mo;
				if(left>=i) {
					pat=pat*combi(left,i)%mo*inv[y+1]%mo;
					left-=i;
				}
			}
		}
	}
	
	cout<<dp[1][0]<<endl;
}

まとめ

かなり手間取ったけど、ここまで１時間で解けたのは良かった。