合成数modにも対応した二項係数を実装する(C++)

〜ステップ1〜

まず、任意の素数modに対応した二項係数について考えましょう。競プロで頻出の998244353や1000000007は素数であり、 $N_i,K_i \leq 200000$ 程度であるため、 $N_i!,K_i!$ は998244353や1000000007で割り切れることがなく、必ずモジュラー逆数がするので、階上とそのmod下でのモジュラー逆数を事前計算すれば計算することができました。しかし、全ての素数となると、例えば $\bmod 2$ で考えると、 $2!$ はモジュラー逆数が存在しません。

<発想>

$\bmod p$ の時、 $N!$ にpで何回わりきれる回数とpで割り切れるだけ割った後のpで割った余りという要素を持たせることを考えます。そうすると、 $N!$ に対応したpで割り切れる回数と割り切れるだけ割った後の余りをそれぞれ、 ${\mathrm{ord}_N,\mathrm{fac}_N}$ として、

$_nC_k≡\mathrm{fac}_n \times (\mathrm{fac}_k \times \mathrm{fac}_{n-r})^{-1} \times p^{(\mathrm{ord}_n-\mathrm{ord}_k-\mathrm{ord}_{n-k})} \pmod p$

となります。ちなみにこれは、素数pに対して $\bmod p^k$ なら上記の $\bmod p$ を $\bmod p^k$ に変えるだけで適応できます。

〜ステップ2〜

$_nC_k \bmod M$ (Mは任意の自然数)について考えます。まず、 $M=p_1^{e_1} \times p_2^{e_2} \times \cdots \times p_k^{e_k}$ と素因数分解できるとします。そうすると、 $_nC_k \bmod p_i^{e_i}$ ( $1 \leq i \leq k$ )はステップ1で考えたようにできます。ここで、 $\mathrm{lcm}(p_1^{e_1},p_2^{e_2}, \cdots ,p_k^{e_k})=M$ なので拡張ユークリッドの互除法を用いると、 $_nC_k \bmod M$ を求めることができます。

〜実装例〜

#include<bits/stdc++.h>
using namespace std;
using ll=long long;
using P=pair<ll,ll>;
// mod function
ll mod(ll a, ll mod) {
    return (a%mod+mod)%mod;
}
ll modpow(ll a, ll n, ll mod) {
    ll res = 1;
    while (n > 0) {
        if (n & 1) res = res * a % mod;
        a = a * a % mod;
        n >>= 1;
    }
    return res;
}
vector<P> prime_factorize(ll N) {
  vector<P> res;
  for (ll a = 2; a * a <= N; ++a) {
    if (N % a != 0) continue;
    ll ex = 0;
    while(N % a == 0){
      ++ex;
      N /= a;
    }
    res.push_back({a, ex});
  }
  if (N != 1) res.push_back({N, 1});
  return res;
}
ll modinv(ll a, ll mod) {
    ll b = mod, u = 1, v = 0;
    while (b) {
        ll t = a/b;
        a -= t * b, swap(a, b);
        u -= t * v, swap(u, v);
    }
    u %= mod;
    if (u < 0) u += mod;
    return u;
}
ll extGcd(ll a, ll b, ll &p, ll &q) {  
    if (b == 0) { p = 1; q = 0; return a; }  
    ll d = extGcd(b, a%b, q, p);  
    q -= a/b * p;  
    return d;  
}
P ChineseRem(const vector<ll> &b, const vector<ll> &m) {
  ll r = 0, M = 1;
  for (int i = 0; i < (int)b.size(); ++i) {
    ll p, q;
    ll d = extGcd(M, m[i], p, q);
    if ((b[i] - r) % d != 0) return make_pair(0, -1);
    ll tmp = (b[i] - r) / d * p % (m[i]/d);
    r += M * tmp;
    M *= m[i]/d;
  }
  return make_pair(mod(r, M), M);
}
struct Comb{
  unordered_map<int,tuple<ll,vector<ll>,vector<ll>>> mp;
  int n_,m;
  ll p_, pm_;
  vector<P> pf;
  Comb(int n,int M) : n_(n) { 
    m=M;
    pf=prime_factorize(M); 
    COMinit(n);
  }
  Comb(ll p, ll pm, int n) :
    n_(n), p_(p), pm_(pm) {
    init(p, pm);
  }
  void init(long long p, long long pm) {
    p_=p,pm_=pm;
    auto&[pms,ord,fac]=mp[p];
    pms=pm;
    ord.resize(n_);
    fac.resize(n_);
    ord[0]=ord[1]=0;
    fac[0]=fac[1]=1;
    for (int i=2;i<n_;i++) {
      long long add=0;
      long long ni=i;
      while (ni % p == 0) add++,ni/=p;
      ord[i]=ord[i-1]+add;
      fac[i]=fac[i-1]*ni%pm;
    }
  }
  void init(long long p, long long pm, int n) {
    init(p, pm);
  }
  void COMinit(int n){
    for(auto p : pf){
      ll ps=p.first,pfs=pow(p.first,p.second);
      init(ps,pfs);
    }
  }
  ll com(ll n, ll r,int p) {
    if (n<0 || r<0 || n<r) return 0;
    auto&[pms,ord,fac]=mp[p];
    ll e=ord[n]-ord[r]-ord[n-r];
    ll res=fac[n]*modinv(fac[r]*fac[n-r]%pms,pms)%pms;
    res=res*modpow(p,e,pms)%pms;
    return res;
  }
  ll operator()(int n, int k){
    if(n<0 || k<0 || n<k) return 0;
    vector<long long> vb, vm;
    for (auto ps : pf) {
        long long p = ps.first, e = ps.second;
        long long pm = pow(p,e);
        long long b = 1;
        b *= com(n, k ,p) % pm;
        b %= pm;
        vm.push_back(pm);
        vb.push_back(b);
    }
    auto res = ChineseRem(vb,vm);
    return res.first;
  }
};
int main(){
  Comb C(200000,998244353);
  cout<<C(5,2)<<'\n';
}