Pythonで逆元を使ってnCr mod 1000000007を計算

AtCoderを始めてから、1000000007で割ったあまりを求めよ、という問題を見る機会があった。

直近だとAtCoder Beginner Contest 156のD問題 Bouquet。

D - Bouquet

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

足し算、引き算、掛け算の場合は計算途中でmodを取ればいいのだが、わり算の場合がややこしい。

次は本番中に解けるよう、nCrと逆元について調べてみた。

nCr

この問題ではrに相当する部分が最大でも10^5であるが、nは最大10^9とかなり大きい。

そのため、nCr = n! / (n-r)! r! と式変形していては間に合わない。

そこで、 nCr = n*(n-1)*(n-2)…(n-k+1) / r*(r-1)*(r-2)…1 と考える。

高校数学で 6C3 = 6*5*4 / 3*2*1 と計算していたのと同じやり方を採用することで、計算時間がO(r)となり間に合う。

分子、分母はそれぞれ掛け算のみで構成されるので、計算途中でmodを取ってOK。

from functools import reduce

numerator = reduce(lambda x, y: x * y % mod, [n - r + k + 1 for k in range(r)])
denominator = reduce(lambda x, y: x * y % mod, [k + 1 for k in range(r)])

逆元

逆元の計算にはフェルマーの小定理を使うが、詳細な理屈は他サイトに譲る。というより僕もよく分かっていない。

ひとまず実装だけでもしたいと思っていた所、ABC34の解説スライドの23ページが大変わかりやすかった。

AtCoder Beginner Contest 034 解説 from AtCoder Inc.

nCr mod 1000000007 の場合、分母 denominator で割り算してmodを取ることと、denominator^(1000000007-2)を掛け算してmodを取ることは等しくなる。

割り算が掛け算に変換され、計算途中でmodを取ることが出来るようになった。

計算量が大幅に減り、mod計算における逆元の有り難さがわかる。

numerator * pow(denominator, mod - 2, mod) % mod

というわけで、提出したコードがこちら。

from functools import reduce

n, a, b = map(int, input().split())
mod = 10 ** 9 + 7


def cmb(r):
    numerator = reduce(lambda x, y: x * y % mod, [n - r + k + 1 for k in range(r)])
    denominator = reduce(lambda x, y: x * y % mod, [k + 1 for k in range(r)])
    return numerator * pow(denominator, mod - 2, mod) % mod


ans = pow(2, n, mod) - 1 - cmb(a) - cmb(b)
print(ans % mod)

参考

「998244353 で割ったあまり」の求め方を総特集！〜逆元から離散対数まで〜 - Qiita

プログラミングコンテストの問題に挑んでいると、「ある量」を求めたい。ただし答えがとても大きくなる可能性があるので、その量をで割ったあまりを求めよ。というタイプの出題を…

フェルマーの小定理の証明と使い方 - Qiita

今回は初等整数論の面白い話題の 1 つである Fermat の小定理を中心に、周辺の話題を紹介していきます。0. はじめに整数論はとても楽しいです。最近では大学受験においても整数論の出題が目立つ…