二項係数 mod 素数のあまり見かけない計算方法

atcoder.jp

 

これを頑張って計算すると以下の式(を998244353で割った余り)を出力すればいいことがわかります。

$$\sum_{M=0}^{N}\binom{N}{M}\cdot \binom{K+M+1}{N+1}$$

この式の導出方法はこの記事の本筋からそれるので省略します。解説参照!

さて、制約によると$K\leq 10^{18},N\leq 3\times 10^5$でmodの法は998244353(素数)みたいです。これを時間内に求めるにはどうすればいいでしょう???

 

うまくいかない方法1

二項係数のmod 素数を求める典型テクニックは、

$1!,2!,...,n!$と、$1/1!,...,1/n!$のテーブルを$O(n)$で作ってmod逆元を組み合わせる方法が一般的である。しかしこの方法の場合、$K+N+1$くらいの大きさのテーブルを作る必要があるが、そんなクソデカテーブルなど作れるわけがないのでこの方法ではうまくいきません。

 

うまくいかない方法2

$$\binom{K+M+2}{N+1}=\dfrac{K+M+2}{K+M+1-N}\binom{K+M+1}{N+1}$$

となるため、$\binom{K+1}{N+1}$を求めてから順々に$\binom{K+M+1}{N+1}$を計算するとうまくいきそうに見えるかもしれない。

しかし、この場合、$K$が998244353よりも大きくなる場合があるので、$K+M+1-N$が998244353の倍数となる可能性がある。そうなるとゼロ除算になってしまうため、この方法だと正常に計算できない。よって、うまくいかない。

 

うまくいかない方法3

judge.yosupo.jp

$m $が大きすぎて間に合いません!!!

結局、Lucasの定理を使おうにも、998244353個分の階乗modを記録する必要があり、それは無理である。

うまくいかない方法4

要は$\binom{K+M+1}{N+1}$を出力するクエリを高速化すればよい。

で、これはつまり

$$(K+M+1)\cdot (K+M)\cdot \cdots \cdot (K+M-N+1)$$

を求めれればよい。

つまり、

$$ a_i=(K - N + 1)\cdot (K - N + 1 + 1)\cdot \cdots \cdot (K - N + i) $$

という順列を前計算で取得して、

$$\dfrac{a_{N+M+1}}{a_{M}}$$

を計算すればよいのでは?

と、思いたくなるが、この場合も、$K-N+1,\ldots,K+N+1$の中に998244353の倍数が混じっていた場合正しく計算できなくなる。よって、うまくいかない。

 

うまくいく方法の例

実は、うまくいかない方法4は実はいい線をいっている。

区間積の方法を累積積で計算したのがまずかった。以下の方法だとうまくいく。

長さ$2N+1$のセグメントツリーを定義する。ここで、

・初期値は$K-N+1,\ldots,K+N+1$という長さ$2N+1$の列

演算子は$x\otimes y=(x*y) \mod 998244353$

単位元は(もちろん)1

とする。

このセグメントツリーに対して、$M $ごとに$a_M\otimes \cdots \otimes a_{N+M}$を計算して出力する。

 

するとクエリ一回あたりの計算量は$O(\log{N})$となり、最終的に全体の計算量は$O(N\log{N})$となるので通る。

また、$K-N+1,\ldots,K+N+1$の間に998244353の場合があっても正しく動く。

セグメントツリーは任意のモノイドに対して動く。

モノイドとは、結合法則を満たして単位元が存在すればよかった。

つまり、逆元が存在しなくても問題ないので、区間内に998244353の倍数が存在しても正しく動くのである。