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. はじめに整数論はとても楽しいです。最近では大学受験においても整数論の出題が目立つようになりましたし、計算機科学においても整...