n! の計算
こいつを速くします。
※ 単純な低レベル処理はcythonやnumbaにより局所的にコンパイルする方が標準的な思想だと思います。また、numpyのufuncを自作するような方法もあります。今回記事にしたものは、atcoder環境下ということを念頭に置いたやや不自然な解決策だとは思います。
※ より一般に、長い1次元配列のmodを取りながらcumprod(累積積?)をとりたい場合に適用可能です。numpy.cumprod あるいは numpy.multiply.accumulate を適用すると、int64・uint64・float64などの積を連続的に適用することになってmodを取ってくれません。
トリッキーですが、これで階乗 mod pの配列が作れています。「積をとりながらmodをとることを繰り返す」という直列処理は簡単に書けないように思いますが、「同時にたくさんの積をとる」「同時にたくさんのmodをとる」という並列処理が簡単に書けることを利用します。
例えば、8! までの計算は、次のような手順で行われています。
【1】適当に数を並べる。(np.arange・reshape)
$1$ | $1$ | $2$ |
$3$ | $4$ | $5$ |
$6$ | $7$ | $8$ |
【2】横にループして積をとる
$0!$ | $1!$ | $2!$ |
$3$ | $3×4$ | $3×4×5$ |
$6$ | $6×7$ | $6×7×8$ |
【3-1】2行目に2!をかける
$0!$ | $1!$ | $2!$ |
$3!$ | $4!$ | $5!$ |
$6$ | $6×7$ | $6×7×8$ |
【3-2】3行目に5!をかける → 完成!
$0!$ | $1!$ | $2!$ |
$3!$ | $4!$ | $5!$ |
$6!$ | $7!$ | $8!$ |
階乗の逆数・二項係数
$p$ を素数とし、$n<p$ とするとき、$n!$ は $\pmod{p}$ で逆数を持ちます。これを計算するには、$(n-1)!^{-1} = (n!)^{-1}\cdot n$ を利用できます。したがって、
\[ ((n-k)!)^{-1} = (n!)^{-1} \times \bigl(n\cdot(n-1)\cdots (n-k+1)\bigr)\] のように、定数に累積積を乗じた形で表すことができます。$n,n-1,n-2,\ldots$の累積積を上述の方法で計算して、それに $n!^{-1}$ を乗ずれば、階乗の逆数が出来上がりです。
固定された $n$ に対して二項係数の列 $\binom{n}{0}, \binom{n}{1},\ldots$ を作るには、
\[ \binom{n}{k} = \frac{n!}{k!(n-k)!} \] を利用します。例えば、$\binom{10000}{k} \pmod{p}$ が入ったarrayを作るには、fact[10000] * fact_inv[:10001] * fact_inv[10000::-1] とすればよいですね。
備考
numpyは、ひとつひとつランダムアクセスすると遅いので、次のようなコードは最悪です。
また計算結果はnp.darrayなので、同じ理由により、作成後の処理によっては階乗の計算後にまたlistに戻した方が有利なこともありえます。この場合は、fact = fact.tolist() などとしてリストに戻せばOKです。tolistって競プロ勢なら喜びそうなメソッドじゃないですか?
個人的には、「ほぼ全部numpy.darrayで計算させたいのに唯一階乗だけうまく書けない!」というところで悩むことが多かったので、今回の発見は割と嬉しいです。