2024.02.26

記事内に商品プロモーションを含む場合があります

PythonでMOD付きの組み合わせ（nCk）を求める

tadanori

競プログラミング（AtCoder）で問題を解いていると、「1000000007で割った余りを求めてください」といった問題が出題されています。

ここでは、組み合わせ数$_nC_k$を求めるプログラムをPythonで書く方法を説明します

結果だけ欲しい方は以下のコードをコピペしてください。使い方は、記事中の「MOD付き組み合わせを計算するクラスの定義」にあります

class cMod :
    def __init__(self, mod, size):
        fact = [1 for _ in range(size)]
        ifact = [1 for _ in range(size)]

        for i in range(1, size) :
            fact[i] = fact[i-1] * i % mod

        ifact[size-1] = pow(fact[size-1], mod-2, mod)

        for i in range(size-2, -1, -1):
            ifact[i] = ifact[i+1] * (i+1) % mod
        self.fact = fact
        self.ifact = ifact


    def nCk(self, n: int, k: int) -> int:
        if n < k or k < 0 : return 0
        return self.fact[n] * self.ifact[k] % mod * self.ifact[n-k]%mod

その他のAtCoderに役立つ記事の一覧は以下にあります。

剰余演算（MOD演算）

剰余演算

剰余演算（modulo演算）は、ある数を別の数で割った際の余りを計算する演算です

aをbで割った場合の余り剰余rは、プログラミングではr = a%bで表されます。

また、以下のように表すこともあります。

$$
a \equiv r \pmod b
$$

正確には、上記の式は「aとrがbを法として同じ剰余を持つ」という意味で厳密には異なりますが、とりあえず上記のイメージで良いと思います。

例えば、10%3は10を3で割った結果で10%3=1となります。

プログラミングではこの剰余演算は「mod計算」などと呼ばれます

AtCoderで剰余演算が使われる理由

AtCoderでは、結果が整数（64bit）を超える場合に剰余演算が使われます。多くのプログラミング言語が64bit整数までしかサポートしていないため、modを使わないとこれを超えるような問題が出題できなくなるためです。

例えば、組み合わせ問題などの場合、$_{100}C_{50}$を計算するだけで、64bitで表現できる範囲を超えてしまいます。

プログラムでの記述

Pythonでは算術演算子%を使ってmod計算を行うことができます。これは、他の多くの言語と共通です。

mod計算の例

>>> 50%5
0
>>> 34%11
1
>>> -10%3
2

豆知識
Pythonでは、浮動小数点型でもmod計算することが可能です。例えば、6%2.5 = 1.0の計算が可能です

Pythonでは、divmodを使うことで、商よ剰余を同時に計算することが可能です。

divmodの例

>>> v, r = divmod(10,3)
>>> print(v, r)
3 1

また、powを使った累乗計算では、剰余を引数とした計算を行うことができます。

powの例

>>> pow(2,4,10)
6

例は$2^4 \pmod {10}$を計算するもので、16%10=6となります。

以上、剰余計算について簡単に説明しました。

先ほど「AtCoderで剰余演算が使われる理由」に書いたように、組み合わせの計算では簡単に64bit整数の範囲を超えてしまいます。このため、AtCoderでは、組み合わせ数を求める場合に剰余演算を行うことが多いです。

以下では、modを含めた組み合わせ計算について解説します。

組み合わせの公式

「n個の中からk個を取り出す組み合わせの場合の数」は、$_nC_k$と表記されます（CはCombination（組合せ）の頭文字です）。

$_nC_k$は、以下の公式で求めることができます。

$$
_nC_k = \frac{n!}{k! (n-k)!}
$$

式から、$n!$, $k!$, $(n-k)!$といった階乗が計算できれば組み合わせは計算できることがわかります。

ただ、$k!$, $(n-k)!$については、$\frac{1}{k!}$, $\frac{1}{(n-k)!}$を求める必要があります。これを求めるために「逆元」の計算が必要となります。

逆元

合同式(mod)の場合、問題となるのが割り算です。

足し算、引き算、掛け算は、そのまま演算して余りを求めれば良いので問題になりません。

割り算については、そのまま計算することができないため工夫が必要になります。

具体的には、割り算の計算を行うには、フェルマーの小定理を利用します。

フェルマーの小定理は、$p$を素数とし、$a$を$p$の倍数でない正数とすると

$$a^{p-1} \equiv 1 \pmod p$$

が成立するというものです。

これを変形すると、以下の式が導き出せます。

$$a^{p-2}\cdot a \equiv 1 \pmod p$$

$a$に$a^{p-2}$を掛け合わせると1になるということは、$a^{p-2} = \frac{1}{a}$ということになります。

つまり、$a^{p-2}$は$a$の逆元となります

これを利用することで、mod pでの$\frac{1}{k!}$, $\frac{1}{(n-k)!}$が計算できます。

MOD付きの組み合わせをクラスとして定義する

AtCoderの問題では、何回も組み合わせを計算する場合があります。

毎回、階乗を計算すると時間がかかります。そこで必要な範囲をあらかじめ計算してテーブル（配列）に保存しておくことにします。組み合わせ計算を何度も行う場合、これで計算を高速化できます。

階乗（n!）の計算

fact[i]に$i!$の結果を格納するとした場合、階乗は以下のコードでn-1まで計算できます。プログラム中のmodは「1000000007で割った余りを求めてください」の場合は、1000000007です。

これで、fact[0]〜fact[n-1]までが計算できます。

fact = [1 for _ in range(n)]
for i in range(1, n):
    fact[i] = fact[i-1] * i % mod

逆元（1/n!）の計算

組み合わせの計算を行う場合、$\frac{1}{k!}$などが必要となります。

Pythonでは、以下のコードで計算することができます。

しかしながら、これをn個のすべてのfact[k]に対して行うと処理時間かかります

powの計算でループがあるので、それだけ時間が必要になります。

pow(k, mod-2, mod)

そこで、少しだけ工夫します。

$$
\frac{1}{(n-1)!} = \frac{1}{1\times 2 \times … \times (n-2) \times(n-1)}
$$

から、

$$
\frac{1}{(n-2)!} = \frac{1}{1\times 2 \times … \times (n-2) \times(n-1)} \times (n-1) = \frac{1}{(n-1)!} \times (n-1)
$$

となります。これを使えば、$\frac{1}{(n-1)!} $だけフェルマーの小定理を使って計算して、あとは、n-1から順番に計算していけば良いことになります。

逆元をifact[k]とすると、以下の計算により前計算することができます。

# fact[n-1]の逆元を計算
ifact[n-1] = pow(fact[n-1], mod-2, mod)
# n-1から順番に計算
for i in range(size-2, -1, -1):
    ifact[i] = ifact[i+1] * (i+1) % mod

MOD付き組み合わせを計算するクラスの定義

これであらかじめ0〜n-1までの階乗が計算できました。これを計算しておけば、あとは、組み合わせの数を計算するだけです。

まとめると、$_nC_k$を求める関数（ここでは、クラスにしました）は以下のようになります。

class cMod :
    def __init__(self, mod, size):
        fact = [1 for _ in range(size)]
        ifact = [1 for _ in range(size)]

        for i in range(1, size) :
            fact[i] = fact[i-1] * i % mod

        ifact[size-1] = pow(fact[size-1], mod-2, mod)

        for i in range(size-2, -1, -1):
            ifact[i] = ifact[i+1] * (i+1) % mod
        self.fact = fact
        self.ifact = ifact


    def nCk(self, n: int, k: int) -> int:
        if n < k or k < 0 : return 0
        return self.fact[n] * self.ifact[k] % mod * self.ifact[n-k]%mod

使い方

まず、以下のようにMODする値と、事前に剰余計算する範囲を引数として初期化を行います。

m = cMod(1000000007, 2*10**5)

１つ目の引数は、modの値（割って余りを求める値）を、２つ目の引数は、あらかじめ用意しておく階乗の配列のサイズです。

初期化後は、以下のようにしてMOD付きの組み合わせが計算できます

m.nCk(10,5)

問題を解いてみる

ここでは、作成した関数を使って、以下の問題を解いてみます。

https://atcoder.jp/contests/math-and-algorithm/tasks/math_and_algorithm_ar

class cMod :
    def __init__(self, mod, size):
        fact = [1 for _ in range(size)]
        ifact = [1 for _ in range(size)]

        for i in range(1, size) :
            fact[i] = fact[i-1] * i % mod

        ifact[size-1] = pow(fact[size-1], mod-2, mod)

        for i in range(size-2, -1, -1):
            ifact[i] = ifact[i+1] * (i+1) % mod
        self.fact = fact
        self.ifact = ifact


    def nCk(self, n: int, k: int) -> int:
        if n < k or k < 0 : return 0
        return self.fact[n] * self.ifact[k] % mod * self.ifact[n-k]%mod
        
    
mod = 1000000007
size = 210000

m = cMod(mod, size)


X, Y = map(int, input().split())
print(m.nCk(X+Y, X))