フェルマーの小定理を用いたCombination(組み合わせ)計算

前書き

AGC025のB問題にて，膨大なCombination計算を強いられたため，フェルマーの小定理を用いることにした．

フェルマーの小定理

　$M$が素数，$k$が任意の自然数で，$M$と$k$が互いに素のとき，

k^{M-1} \equiv 1 \pmod{M}

となる．[証明はこちら]
　これを変形して，

k^{M-1} = k \cdot k^{M-2}  \equiv 1 \pmod{M}

ゆえに

k^{-1} \equiv k^{M-2} \pmod{M}

が導かれる．

Combination計算への応用

{}_n C _k = \cfrac{n!}{k!\cdot(n-k)!}=n!\cdot (k!)^{-1}\cdot ((n-k)!)^{-1}

と書ける．
上式($k^{-1} \equiv k^{M-2} \pmod{M}$)と組み合わせると，

{}_n C_k\equiv n!\cdot (k!)^{M-2} \cdot ((n-k)!)^{M-2} \pmod{M} 

となる．以下ではこれを実装する．

実装

　実際にやってみた．以下に${}_n C_k$の計算を行うコードを示す．ただし，$0< n \leqq 300000$であるとする．

test.cpp

#include <iostream>
#include <vector>

using namespace std;
typedef long long ll;

const ll M = 998244353;

vector<ll> fac(300001); //n!(mod M)
vector<ll> ifac(300001); //k!^{M-2} (mod M)
//a,bの範囲的にこれだけ配列を用意していけば十分

ll mpow(ll x, ll n){ //x^n(mod M) ←普通にpow(x,n)では溢れてしまうため，随時mod計算
    ll ans = 1;
    while(n != 0){
        if(n&1) ans = ans*x % M;
        x = x*x % M;
        n = n >> 1;
    }
    return ans;
}

ll comb(ll a, ll b){ //aCbをmod計算
    if(a == 0 && b == 0)return 1;
    if(a < b || a < 0)return 0;
    ll tmp = ifac[a-b]* ifac[b] % M;
    return tmp * fac[a] % M;
}

int main()
{
    ll n,k;
    cin >> n >> k;

    //大した量ではないので，先にfax[i]とifax[i]を全て計算しておく
    fac[0] = 1;
    ifac[0] = 1;
    for(ll i = 0; i<300000; i++){
        fac[i+1] = fac[i]*(i+1) % M; // n!(mod M)
        ifac[i+1] = ifac[i]*mpow(i+1, M-2) % M; // k!^{M-2} (mod M) ←累乗にmpowを採用
    }


    ll ans=0;
    ans = comb(n, k)%M;
    cout << ans << endl;
    return 0;
}

応用

実際の問題に応用してみた．問題はyukicoder No.117を用いる．当然フェルマーの小定理は順列や重複組み合わせにも応用可能．

順列計算への応用

{}_n P _k = \cfrac{n!}{(n-k)!}=n!\cdot ((n-k)!)^{-1}

と書ける．
上式($k^{-1} \equiv k^{M-2} \pmod{M}$)と組み合わせると，

{}_n P_k\equiv n!\cdot ((n-k)!)^{M-2} \pmod{M} 

となる．

重複組み合わせ計算への応用

{}_n H _k = \cfrac{(n+k-1)!}{k!\cdot(n-1)!}=(n+k-1)!\cdot (k!)^{-1}\cdot ((n-1)!)^{-1}

と書ける．
上式($k^{-1} \equiv k^{M-2} \pmod{M}$)と組み合わせると，

{}_n H_k\equiv (n+k-1)!\cdot (k!)^{M-2} \cdot ((n-1)!)^{M-2} \pmod{M} 

となる．

main.cpp

#include <iostream>
#include <vector>
#include <cmath>

#define FOR(i,a,b) for(int i=(a);i<(b);++i)

using namespace std;
typedef long long ll;

const ll M = pow(10,9)+7;

vector<ll> fac(1000001); //n!(mod M)
vector<ll> ifac(1000001); //k!^{M-2} (mod M)


ll mpow(ll x, ll n){ //x^n(mod M)
    ll ans = 1;
    while(n != 0){
        if(n&1) ans = ans*x % M;
        x = x*x % M;
        n = n >> 1;
    }
    return ans;
}

ll comb(ll a, ll b){ //aCb(mod M)
    if(a == 0 && b == 0)return 1;
    if(a < b || a < 0)return 0;
    ll tmp = ifac[a-b]* ifac[b] % M;
    return tmp * fac[a] % M;
}

ll perm(ll a, ll b){
    if(a < b) return 0;
    ll tmp = ifac[a-b] % M;
    return tmp * fac[a] % M;
}

ll dupc(ll a,ll b){
    if(a == 0 && b == 0)return 1;
    if(a < 0 || b == 0)return 0;
    ll tmp = ifac[a-1]* ifac[b] % M;
    return tmp * fac[a+b-1] % M;
}

int main()
{
    ll t;
    cin >> t;
    vector<char> c(t);
    vector<int> a(t),b(t);
    FOR(i, 0, t){
        scanf(" %c(%d,%d)",&c[i] ,&a[i] ,&b[i]);
    }

    fac[0] = 1;
    ifac[0] = 1;
    for(ll i = 0; i<1000000; i++){
        fac[i+1] = fac[i]*(i+1) % M; // n!(mod M)
        ifac[i+1] = ifac[i]*mpow(i+1, M-2) % M; // k!^{M-2} (mod M)
    }

    FOR(i, 0, t){
        ll ans=0;
        if(c[i] == 'P')ans = perm(a[i],b[i]);
        if(c[i] == 'C')ans = comb(a[i],b[i]);
        if(c[i] == 'H')ans = dupc(a[i],b[i]);
        cout << ans << endl;
    }


    return 0;
}

あとがき

とりあえず，膨大なCombination(組み合わせ)計算を強いられたときは，素数$M$についてmodをとるフェルマーの小定理を用いて計算すべしということが分かった．

参考文献

マツシタのお勉強：http://keita-matsushita.hatenablog.com/entry/2016/12/05/184011
高校数学の美しい物語：https://mathtrain.jp/fermat_petit