前提
ABC156 D Bouquet を解くために必要になった。この問題の条件は
$p = 10^{9}+7$
$2 \leq n \leq 10^{9}$
$1 \leq r \leq 2 \times 10^{5}$
後に ABC167 E Colorful Blocks を解くときに、拡張されたユークリッドの互除法を使ったので追加した。
結論
フェルマーの小定理を使った方法
def ncr(n, r, mod):
ret = 1
for i in range(1, r+1):
ret = (ret * (n-i+1) * pow(i, mod-2, mod)) % mod
return ret
拡張ユークリッドの互除法を使った方法
def ncr_eu(n, r, mod):
ret = 1
if r < n:
inv = [1]
for i in range(1, r+1):
inv.append(max(1, (-(mod//i) * inv[mod % i]) % mod))
ret = ret*(n+1-i)*inv[i] % mod
return ret
ただし mod は素数
解説
組み合わせの数
$n$ 個の中から $r$ 個を選ぶ組み合わせは
_n C _r = \frac{n!}{r!(n-r)!}
なので単純に
from math import factorial
def ncr(n, r):
return factorial(n) // factorial(r) // factorial(n-r)
としたいところだが、非常に遅い。
フェルマーの小定理
早く計算するために $x!\pmod{p}$ を計算して factorial を置き換えたいが、単純に
_n C _r \equiv \frac{n!\bmod p}{(r!\bmod p)((n-r)!\bmod p)}\bmod p
とはならないらしい。そこで、フェルマーの小定理を持ち出す。
a^{p-1} \equiv 1 \bmod p
$a$ と $p$ が互いに素なとき、両辺を $a$ で割ることができ、
a^{p-2} \equiv a^{-1} \bmod p
が得られる。 この $a^{-1}$ を逆元（ぎゃくげん）というらしい。
これを使うと
_n C _r \equiv (n!\bmod p)((r!)^{p-2}\bmod p)(((n-r)!)^{p-2}\bmod p) \bmod p
となって、掛け算だけであらわせる。
繰り返し二乗法
ここで多用する $x^{p-2}$ を高速に計算するために、
\begin{eqnarray}
x^{2} &=& x \cdot x \\
x^{4} &=& x^{2} \cdot x^{2} \\
x^{8} &=& x^{4} \cdot x^{4} \\
\cdots
\end{eqnarray}
として計算量を減らすテクニックだが、 Python では標準の pow を使ったほうが早い。
実装例
def bpow(x, y, z):
a = 1
while y:
if y & 1:
a = (a*x) % z
x = (x*x) % z
y >>= 1
return a
拡張されたユークリッドの互除法
ユークリッドの互除法といいつつ、もっと直感的な方法で求められる。$p$ を $a$ で割った商を $q$、剰余を $r$ とすると、
\begin{eqnarray}
p &=& qa+r \\
0 &\equiv& qa + r \bmod p \\
0 &\equiv& q + a^{-1} r \bmod p \\
a^{-1} &\equiv& -q \cdot r^{-1} \bmod p \\
\end{eqnarray}
区間 $[1, a)$ の逆元を使って$a$ の逆元を計算するので、$_n C _r$ の $n$ が固定で $r$ が複数必要になる場合に効率が良い。
参考：よくやる二項係数 (nCk mod. p)、逆元 (a^-1 mod. p) の求め方 - けんちょんの競プロ精進記録