FPS24 C - 数列

考え方

分数型の母関数の扱いが登場。

数列の $1$ つの要素に $M$ 以下の数を割り当てることを表現する母関数は、以下。
$$
A(x)=1+x+x^2+x^3+\dots +x^M
$$

これを $N$ 乗して、$x^S$ の係数を答えればよい。
求める方針を $2$ つ紹介する。

方針 $1$

人力で畳み込みをしやすくするための式変形を頑張る。
まず、$A(x)$ を以下のように変形する。
$$
A(x)=1+x+x^2+x^3+\dots +x^M=\dfrac{1-x^{M+1}}{1-x}
$$

これを $N$ 乗すると、以下のようになる。
$$
\{A(x)\}^N=\dfrac{(1-x^{M+1})^N}{(1-x)^N}
$$

これの分子側は、普通に二項展開すればよい。
$$
(1-x^{M+1})^N = ({}_N\mathrm{C}_0-{}_N\mathrm{C}_1x^{M+1}+{}_N\mathrm{C}_2x^{2M+2}-{}_N\mathrm{C}_3x^{3M+3}+\dots+(-1)^{N}{}_N\mathrm{C}_Nx^{NM+N})
$$

分母側は、以下の公式で $r=N$ とすることで、級数の形に戻すことができる。
$$
\dfrac{1}{(1-x)^r} = {}_{r-1}\mathrm{C}_{r-1}x^0+{}_{r}\mathrm{C}_{r-1}x^1+{}_{r+1}\mathrm{C}_{r-1}x^2+{}_{r+2}\mathrm{C}_{r-1}x^3+\dots
$$

最後の畳み込みは、素直に指数の和が $S$ になる組を拾って、係数の積の総和を求めればよい。

コード部分は二項係数ライブラリさえあればシンプルで、階乗の事前計算分に加えて $O(S)$ である。

方針 $2$

少し発展的な知識を使って力技で求める。
つまり、NTTによる畳み込みを用いて、繰り返し二乗法で $N$ 乗を計算する。

$x^{S+1}$ 以上は意味がないのでそこから先は途中で切り捨てながら求める。
計算量は $O(S\log S\log N)$ となる。
「解答例2」参照。

入力例1での動作

入力を受け取る。

n: 3
m: 2
s: 4

$i=0$ では、分母側の冪は $j=4$ となる。
加える値は ${}_3\mathrm{C}_0\times{}_6\mathrm{C}_2=15$ なので、ここまでの総和は $15$ となる。

$i=1$ では、分母側の冪は $j=1$ となる。
二項係数の積は ${}_3\mathrm{C}_1\times{}_3\mathrm{C}_2=9$ だが、$i$ が奇数なので符号を反転させる。
よって、ここまでの総和は $15-9=6$ となる。

$i=2$ では $i(M+1)=6>S=4$ となるので、これ以上調べる必要はない。
以上より、答えは $6$ となる。

注意点

二項係数は $998244353$ で割った余りを求めながら計算する。
階乗を用いて二項係数を求める際の除算には逆元を用い、逆元は繰り返し二乗法で求める。

別解

特になし。