https://kanpurin.hatenablog.com/entry/2021/09/05/014756

問題概要

${1}$ 以上 ${N}$ 以下の ${k}$ について以下の問題を解け．

${1}$ から ${N}$ まで番号が付いた人を空でない ${k}$ 個の集合に分ける． ${M}$ で割った余りが等しい人は同じ集合に含まれないような分け方の総数を求めよ．

${2\leq N\leq 5000\\ 2\leq M\leq N}$

解説

包除原理

集合も区別して考え，最後に ${k!}$ で割ります．

${M}$ で割った余りが ${i}$ であるような番号の人を ${k}$ 個の集合に分けます．このとき，1つの集合には高々1人含まれる(1人もいない集合があってもいい)．この分け方の数は， ${n_i}$ を ${M}$ で割った余りが ${i}$ であるような番号の人の数として， ${{}_{k}{\rm{P}}_{n_i}}$ となります．

すべての ${0\leq i < M}$ に対してこの値の積を求めると，満たすべき条件のうち「集合は空でない」以外の条件を満たすような分け方の数が分かります．以下，「分け方」を「集合は空でない」以外の条件を満たすような分け方のこととします．

後は包除原理で解けます．

${X_i}$ を「 ${i}$ 番目の集合が空でないような分け方」の集合とします．
${|X_1\cap X_2\cap\cdots\cap X_k|}$ が求めるものです．

${|X_1\cap X_2\cap\cdots\cap X_k|\\ =\left\lvert\overline{\overline{X_1}\cup \overline{X_2}\cup\cdots\cup \overline{X_k}}\right\rvert\\ =\displaystyle\sum_{T\subseteq \{1,2,\cdots ,k\}}(-1)^{|T|}\left|\bigcap_{i\in T}\overline{X_i}\right|}$

${|T|}$ が同じなら ${\displaystyle\left|\bigcap_{i\in T}\overline{X_i}\right|}$ の値は同じであり， ${|T|=m}$ のとき， ${\displaystyle\left|\bigcap_{i\in T}\overline{X_i}\right|=\prod_{i=0}^{M-1}{}_{k-m}{\rm{P}}_{n_i}}$ であるので

${=\displaystyle\sum_{m=0}^{k}{}_{k}{\rm{C}}_{m}(-1)^{m}\prod_{i=0}^{M-1}{}_{k-m}{\rm{P}}_{n_i}}$ となります．

${k!}$ で割ったものが答えとなります． ${\displaystyle\prod_{i=0}^{M-1}{}_{k-m}{\rm{P}}_{n_i}}$ を前計算しておくと ${O(N^2)}$ ．

qiita.com

さらなる高速化

${\displaystyle\frac{1}{k!}\sum_{m=0}^{k}{}_{k}{\rm{C}}_{m}(-1)^{m}\prod_{i=0}^{M-1}{}_{k-m}{\rm{P}}_{n_i}\\ =\displaystyle\frac{1}{k!}\sum_{m=0}^{k}\frac{k!}{m!(k-m)!}(-1)^{m}\prod_{i=0}^{M-1}\frac{(k-m)!}{(k-m-n_i)!}\\ =\displaystyle\sum_{m=0}^{k}\frac{(-1)^{m}(k-m)!^{M}}{m!(k-m)!}\prod_{i=0}^{M-1}\frac{1}{(k-m-n_i)!}}$
$f(x)=\prod_{i=0}^{M-1}\frac{1}{(x-n_i)!}$ とすると
${=\displaystyle\sum_{m=0}^{k}\frac{(-1)^{m}(k-m)!^{M-1}}{m!}f(k-m)\\ =\displaystyle\sum_{s+t=k}\left(s!^{M-1}f(s)\right)\left(\frac{(-1)^t}{t!}\right)}$

となります．
${ n_i= \left\{ \begin{array}{ll} \left\lfloor\frac{N}{M}\right\rfloor +1 & (1\leq i\leq N\% M) \\ \left\lfloor\frac{N}{M}\right\rfloor & (otherwise) \end{array} \right. }$
であるので， ${f(x)}$ は ${O(\log{M})}$ で求まります．
したがって，高速な畳み込みを行えば ${O(N\log{N})}$ で解けます．