【python】nCkを導出する方法(ABC156-D)

ABC156-D にて必要になったので備忘。

何も考えずに実装すると

nC_k = \frac{n!}{(n-k)!k!}

$nC_k$の定義は↑なので、これをそのままコードに書くと、

import math
def comb(n,k):
    math.factorial(n) // (math.factorial(n - k) * math.factorial(k))

print(comb(4,2))
# 6

こんな感じになる。
しかし$n!$は$O(n)$なので、今回のABC156-Dのような$n\leqq10^9$の場合では計算時間が長くなり間に合わない。なんとかして計算時間を短くしなければならない。

前提知識

繰り返し二乗法とフェルマーの小定理をまずは知っている必要がある。

繰り返し二乗法

べき乗を高速に求める方法。
pythonではpow(x, n, mod)とすることで $x^n$ を $mod$ で割ったあまりを計算できる。

フェルマーの小定理

$p$ を素数とし、$a$ を $p$ の倍数でない整数（$a$ と $p$ は互いに素）とするときに、以下が成り立つ。

a^{p-1} \equiv 1 (mod\, p)

要は $a^{p-1}$を$p$で割ったらあまりは1だよと言っている。

nCkを式変形する

$nC_k$は以下のように変形できる。

nC_k = \frac{n!}{(n-k)!k!}=\frac{n(n-1)(n-2)\cdots(n-k+1)}{k!}

ここで$\frac{1}{k!}$に注目する。

フェルマーの小定理から、

a^{p-1} mod\, p = 1 \\ \\
a^{p-2} mod\, p = \frac{1}{a} \\

$a^{p-2} mod\, p = \frac{1}{a}$が得られる。
よって $\frac{1}{k!} = k!^{p-2} mod\, p$が成り立ち、$k!^{p-2}$は繰返し二乗法で高速に求めることができるので、$\frac{1}{k!}$を高速に求めることが可能になる。

以上から$nC_k$は以下で表すことができる。

nC_k = n(n-1)\cdots(n-k+1) \times(k!^{p-2}\,mod\,p)

計算量は$O(k)$である。（$k!$が一番時間がかかる）

実装

# O(k)
def comb(n,k):
    nCk = 1
    MOD = 10**9+7

    for i in range(n-k+1, n+1):
        nCk *= i
        nCk %= MOD

    for i in range(1,k+1):
        nCk *= pow(i,MOD-2,MOD)
        nCk %= MOD
    return nCk

print(comb(4,2))
# 6

参考

https://ja.wikipedia.org/wiki/%E3%83%95%E3%82%A7%E3%83%AB%E3%83%9E%E3%83%BC%E3%81%AE%E5%B0%8F%E5%AE%9A%E7%90%86
https://note.com/tanon_cp/n/ncff944647054