AtCoder ABC #217 : G - Groups

ちょっと計算量心配だったけど問題なかった。
https://atcoder.jp/contests/abc217/tasks/abc217_g

問題

それぞれ1～N番の番号を人をk個のグループに分ける。
ただし、番号がmod Mで一致する人は同じグループに入れられない。
k=1～Nの時それぞれ、グループの分け方は何通りか。

解法

Editorialとは異なる解法。
f(m,g) := mod Mがm未満の人までグループわけを決めたとき、gグループに分かれている組み合わせ

とする。
f(m,g)がわかっているとき、mod Mがmの人がn人いたときのことを考える。
n人のうちa人が既存のgグループのどこかに入るとすると、
f(m+1,g+(n-a)) += f(m,g)*binom(g,a)*P(n,a)
と遷移する。aは0～nまで総当たりすると、この処理は三重ループで一見怖いが実質O(M*N*(N/M))=O(N^2)で済む。

int N,M;
const ll mo=998244353;
int num[55555];

ll from[15050];
ll to[15050];

const int NUM_=400001;
ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];

ll comb(ll N_, ll C_) {
	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>>M;
	FOR(i,N) num[i%M]++;
	from[0]=1;
	FOR(i,M) {
		ZERO(to);
		FOR(j,N+1) if(from[j]) {
			FOR(x,min(j,num[i])+1) {
				(to[j+num[i]-x]+=from[j]*comb(j,x)%mo*comb(num[i],x)%mo*fact[x])%=mo;
			}
		}
		swap(from,to);
	}
	FOR(i,N) cout<<from[i+1]<<endl;
	
}

まとめ

最初mod 998244343で求めようとしてしまい、混乱してしまった。