izumo’s diary

主に競プロの精進記録

二項係数 (nCr) の計算方法

AtCoder Grand Contest 025 にて _n \mathrm{C} _r\ \bmod\ p を使う問題が出題されました。いい機会だと思ったので i!,\ \left(i!\right)^{-1},\ _n \mathrm{C} _r\ \bmod\ p の高速な実装をしてみました。

使用言語: C++

1. なぜ逆元 \left(i!\right)^{-1} が必要なのか

_n \mathrm{C} _r=\frac{n!}{r!\times(n-r)!}
なので、n!,\ r!,\ (n-r)! が分かればO(1)_n \mathrm{C} _r が求まります。よって事前に階乗を計算しておけばよいことが分かります。

しかし残念ながら
_n \mathrm{C} _r\equiv\frac{n!\ \bmod\ p}{(r!\ \bmod\ p)\times((n-r)!\ \bmod\ p)}\ \bmod\ p
とは限りません。

例:p=7 のとき
_6 \mathrm{C} _3\equiv 20\equiv 6\ \bmod\ 7
です。一方、
6!\equiv{720}\equiv{6},\ 3!\equiv{6}\ \bmod\ 7
なので \frac{6!}{3!\times{3!}} は計算できません。

よって実際には \left(i!\right)^{-1}\ \bmod\ p も事前に計算しておく必要があります。

2. 階乗 i! と 逆元 \left(i!\right)^{-1} の計算

i!\equiv(i-1)!\times{i}\ \bmod\ p
であることを使えば \ 1!,\ 2!,…,\ n! は全体で O(n) で求まります。

一方、
\left(\left(i-1\right)!\right)^{-1}\equiv \left(i!\right)^{-1}\times{i}\ \bmod\ p
であるので \left(n!\right)^{-1} が分かれば \left(\left(n-1\right)!\right)^{-1},\ \left(\left(n-2\right)!\right)^{-1},…,\ \left(0!\right)^{-1} は全体で O(n) で求まります。しかし \left(n!\right)^{-1} を求めるのは少々やっかいです。まず思いつくのは逆元の定義
\left(n!\right)^{-1}\times{n!}\equiv1\ \bmod\ p
を利用して愚直に全探索する方法ですが、計算量は O(p) になってしまいます。実は次に紹介するフェルマーの小定理と繰り返し二乗法を利用すると \left(n!\right)^{-1}O(log\ p) で求めることができます。

3. フェルマーの小定理

フェルマーの小定理p素数としたとき p の倍数でない任意の整数 a に対して
a^{p-1}\equiv 1\ \bmod\ p
が成立するというものです。簡単に証明を書いておきます。

(証明)
ap は互いに素なので
i\times{a}\equiv{j\times{a}}\ \bmod\ p
となるi,\ j\ (0<{i,\ j}<{p},\ i\neq{j}) は存在しない。よって a,\ 2a,…,\ (p-1)a\ \bmod\ p はすべて異なり 1,\ 2,…,\ p-1 のいずれかと同じになる。それぞれの要素の積をとると
(p-1)!\times{a^{p-1}}\equiv{(p-1)!}\ \bmod\ p
となるので定理が従う。
(証明終)

この定理を利用すると
\left(n!\right)^{-1}\equiv{\left(n!\right)^{p-2}}\ \bmod\ p
であることが分かります。

4. 繰り返し二乗法

繰り返し二乗法はその名の通り二乗を繰り返して a^n を計算するアルゴリズムです。n の二進法表現を
n=(n_{k-1},\ n_{k-2},…, n_0)_2
とすると、
a^n=\left(a^{2^0}\right)^{n_0}\times{\left(a^{2^1}\right)^{n_{1}}}\times{…}\times{\left(a^{2^{k-1}}\right)^{n_{k-1}}}
と表せます。
a^{2^i}=\left(a^{2^{i-1}}\right)^2
を使って a^{2^1},\ …,\ a^{2^{k-1}} をそれぞれ求めることができるので、計算量は O(\log{} n) となります。

例:a=3,\ n=5=(1,\ 0,\ 1)_2 のとき
\begin{eqnarray}3^5&=& \left(3^{2^0}\right)^1\times{\left(3^{2^1}\right)^0}\times{\left(3^{2^2}\right)^1}\\ &=& 3^1\times{\left(3^2\right)^0}\times{\left(3^{2^2}\right)^1}\\ &=& 3\times{9^0}\times{\left(9^2\right)^1}\\ &=& 243\end{eqnarray}

繰り返し二乗法を利用すると \left(n!\right)^{-1}\equiv \left(n!\right)^{p-2}\ \bmod\ pO(\log{} p) で求まります。

5. 実装

以上より i!\ \bmod\ p\ O(n),\ \left(i!\right)^{-1}\ \bmod\ pO(n+\log{} p) で前計算できることが分かりました。よって _n \mathrm{C} _r\ \bmod\ pO(n+\log{} p) で求まります。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define REP(i, n) for(int i=0; i<(n); ++i)
#define FOR(i, a, b) for(int i=(a); i<(b); ++i)
#define FORR(i, a, b) for(int i=(b)-1; i>=(a); --i)
constexpr ll MOD=1000000007ll;
// constexpr ll MOD=998244353ll;
constexpr int MAX=510000;

ll fact[MAX], fact_inv[MAX];

// 繰り返し二乗法
ll power(ll a, ll b){
	ll res=1;
	while(b>0){
		if(b&1) res=res*a%MOD;
		a=a*a%MOD;
		b>>=1;
	}
	return res;
}

ll comb(ll n, ll r){
	return (fact[n]*fact_inv[r])%MOD*fact_inv[n-r]%MOD;
}

int main(){
	ios::sync_with_stdio(false);
	cin.tie(0);
	int n=100000;
	// cin>>n;
	fact[0]=1;
        // 階乗の計算
	REP(i, n) fact[i+1]=fact[i]*(i+1)%MOD;
	fact_inv[n]=power(fact[n], MOD-2);
        // 逆元の計算
	FORR(i, 0, n) fact_inv[i]=fact_inv[i+1]*(i+1)%MOD;
        // nCrの計算
        cout<<"6C3: "<<comb(6, 3)<<'\n';
        cout<<"666C390: "<<comb(666, 390)<<'\n';
        return 0;
}