二項係数を10^9+7 で割った余りを求める方法

注意 2020年2月23日 以下に書いてある方法はやや遅く、TLEになる危険性もあります。新しく書いたこの記事も参考にしてください。 linus-mk.hatenablog.com


競技プログラミングでよくある「二項係数 nCk を109+7 で割った余りを求める」方法を整理しておく。
Python 3.7.1で書いていますが、(AtCoderで使われている)Python 3.4.3でも動きます。

# 二項係数を 10^9+7 で割った余りを求める
# https://qiita.com/drken/items/3b4fdf0a78e7a138cd9a#5-%E4%BA%8C%E9%A0%85%E4%BF%82%E6%95%B0-ncr
# https://qiita.com/Yaruki00/items/fd1fc269ff7fe40d09a6
# https://www.hamayanhamayan.com/entry/2018/06/06/210256

mod = 10**9 + 7

# x ** a をmodで割った余りを、O(log(a))時間で求める。
def power(x, a):
    if a == 0:
        return 1
    elif a == 1:
        return x
    elif a % 2 == 0:
        return power(x, a//2) **2 % mod
    else:
        return power(x, a//2) **2 * x % mod

# xの逆元を求める。フェルマーの小定理より、 x の逆元は x ^ (mod - 2) に等しい。計算時間はO(log(mod))程度。
# https://qiita.com/Yaruki00/items/fd1fc269ff7fe40d09a6
def modinv(x):
    return power(x, mod-2)

# 二項係数の左側の数字の最大値を max_len とする。nとかだと他の変数と被りそうなので。
# factori_table = [1, 1, 2, 6, 24, 120, ...] 要は factori_table[n] = n!
# 計算時間はO(max_len * log(mod))

max_len = 30 #適宜変更する

factori_table = [1] * (max_len + 1)
factori_inv_table = [1] * (max_len + 1)
for i in range(1, max_len + 1):
    factori_table[i] = factori_table[i-1] * (i) % mod
    factori_inv_table[i] = modinv(factori_table[i])

def binomial_coefficients(n, k):
    # n! / (k! * (n-k)! )
    return factori_table[n] * factori_inv_table[k] * factori_inv_table[n-k]

コメント中に書いたとおり、最初にテーブルを構築するのにO(max_len * log(mod))程度の時間がかかる。

AtCoder Beginner Contest 132 D - Blue and Red Balls を解くために書いたけど、解説見たら「最大2000なので、パスカルの三角形を上から求めていく要領でDP(動的計画法)でも行けますね」って書いてあった。「パスカルの三角形を上から求めていくDP」だと計算量はO(max_len2)になる。 2000×2000のテーブルを埋めるのがPythonだとTLEになりそうで、速いほうの解法にしてしまった。

直近でこの解法を使う必要があったのは、M-SOLUTIONS プロコンオープン C - Best-of-(2n-1) の問題。これは400点だけど、高速解法が必要。400点あたりからは高速な解法が要求されると思っておいたほうが良さそうだ。