Combination （${}_n \mathrm{C}_r \bmod{m}$，Lucas の定理の拡張）
(number/combination.hpp)

• View this file on GitHub
• Last update: 2022-11-15 00:34:03+09:00
• Include: #include "number/combination.hpp"
• Link: https://ferin-tech.hatenablog.com/entry/2018/01/17/010829
• Link: https://judge.yosupo.jp/problem/binomial_coefficient

固定された正整数 $m = p_1^{q_1} \dots p_{k}^{q_k} ( \le 10^7)$ について，$\binom{n}{r} \bmod{m}$ の値を計算する．前計算 $O(\sum_i p_i^{q_i})$，クエリ $O(k \mathrm{log}(n))$．

原理

$m$ を素因数分解し，各素数冪 $p^q$ を法とした $\binom{n}{r}$ の値を拡張 Lucas の定理を用いて計算し，中国剰余定理を用いて解を復元する．

使用方法

コンストラクタには法 $m$ を素因数分解した結果の（素数，次数）の組として（例：vector<pair<int, int>>, map<int, int> ）与える．例えば $\bmod{60}$ で計算したい場合は {2, 2}, {3, 1}, {5, 1} を与えればよい．また，本ライブラリの Sieve クラスの factorize() が返す結果をそのまま渡してもよい：

combination nCr(Sieve(1 << 20).factorize(mod));
cout << nCr(n, r) << '\n';

問題例

• Library Checker: Binomial Coefficient

リンク

• 2015 ICL, Finals, Div. 1　J. Ceizenpok’s formula (nCk mod m の求め方) - ferinの競プロ帳

Depends on

• number/bare_mod_algebra.hpp

Verified with

• number/test/combination.stress.test.cpp
• number/test/combination.test.cpp

Code

#pragma once
#include "bare_mod_algebra.hpp"
#include <utility>
#include <vector>

// nCr mod m = p^q (p: prime, q >= 1)
// Can be used for n, r <= 1e18, m <= 1e7
// Complexity: O(m) (construction), O(log(n)) (per query)
// https://ferin-tech.hatenablog.com/entry/2018/01/17/010829
struct combination_prime_pow {
    int p, q, m;
    std::vector<int> fac, invfac, ppow;

    long long _ej(long long n) const {
        long long ret = 0;
        while (n) ret += n, n /= p;
        return ret;
    }

    combination_prime_pow(int p_, int q_) : p(p_), q(q_), m(1), ppow{1} {
        for (int t = 0; t < q; ++t) m *= p, ppow.push_back(m);
        fac.assign(m, 1);
        invfac.assign(m, 1);
        for (int i = 1; i < m; ++i) fac[i] = (long long)fac[i - 1] * (i % p ? i : 1) % m;
        invfac[m - 1] = fac[m - 1]; // Same as Wilson's theorem
        assert(1LL * fac.back() * invfac.back() % m == 1);
        for (int i = m - 1; i; --i) invfac[i - 1] = (long long)invfac[i] * (i % p ? i : 1) % m;
    }

    int nCr(long long n, long long r) const {
        if (r < 0 or n < r) return 0;
        if (p == 2 and q == 1) return !((~n) & r); // Lucas
        long long k = n - r;
        long long e0 = _ej(n / p) - _ej(r / p) - _ej(k / p);
        if (e0 >= q) return 0;

        long long ret = ppow[e0];
        if (q == 1) { // Lucas
            while (n) {
                ret = __int128(ret) * fac[n % p] * invfac[r % p] * invfac[k % p] % p;
                n /= p, r /= p, k /= p;
            }
            return (int)ret;
        } else {
            if ((p > 2 or q < 3) and (_ej(n / m) - _ej(r / m) - _ej(k / m)) & 1) ret = m - ret;
            while (n) {
                ret = __int128(ret) * fac[n % m] * invfac[r % m] * invfac[k % m] % m;
                n /= p, r /= p, k /= p;
            }
            return (int)ret;
        }
    }
};

// nCr mod m
// Complexity: O(m) space worst (construction), O(log(n) log(m)) (per query)
// Input: pairs of (prime, degree), such as vector<pair<int, int>> and map<int, int>
// https://judge.yosupo.jp/problem/binomial_coefficient
struct combination {
    std::vector<combination_prime_pow> cpps;
    std::vector<int> ms;

    template <class Map> combination(const Map &p2deg) {
        for (auto f : p2deg) {
            cpps.push_back(combination_prime_pow(f.first, f.second));
            ms.push_back(cpps.back().m);
        }
    }

    int operator()(long long n, long long r) const {
        if (r < 0 or n < r) return 0;
        std::vector<int> rs;
        for (const auto &cpp : cpps) rs.push_back(cpp.nCr(n, r));
        return crt(rs, ms).first;
    }
};

#line 2 "number/bare_mod_algebra.hpp"
#include <algorithm>
#include <cassert>
#include <tuple>
#include <utility>
#include <vector>

// CUT begin
// Solve ax+by=gcd(a, b)
template <class Int> Int extgcd(Int a, Int b, Int &x, Int &y) {
    Int d = a;
    if (b != 0) {
        d = extgcd(b, a % b, y, x), y -= (a / b) * x;
    } else {
        x = 1, y = 0;
    }
    return d;
}
// Calculate a^(-1) (MOD m) s if gcd(a, m) == 1
// Calculate x s.t. ax == gcd(a, m) MOD m
template <class Int> Int mod_inverse(Int a, Int m) {
    Int x, y;
    extgcd<Int>(a, m, x, y);
    x %= m;
    return x + (x < 0) * m;
}

// Require: 1 <= b
// return: (g, x) s.t. g = gcd(a, b), xa = g MOD b, 0 <= x < b/g
template <class Int> /* constexpr */ std::pair<Int, Int> inv_gcd(Int a, Int b) {
    a %= b;
    if (a < 0) a += b;
    if (a == 0) return {b, 0};
    Int s = b, t = a, m0 = 0, m1 = 1;
    while (t) {
        Int u = s / t;
        s -= t * u, m0 -= m1 * u;
        auto tmp = s;
        s = t, t = tmp, tmp = m0, m0 = m1, m1 = tmp;
    }
    if (m0 < 0) m0 += b / s;
    return {s, m0};
}

template <class Int>
/* constexpr */ std::pair<Int, Int> crt(const std::vector<Int> &r, const std::vector<Int> &m) {
    assert(r.size() == m.size());
    int n = int(r.size());
    // Contracts: 0 <= r0 < m0
    Int r0 = 0, m0 = 1;
    for (int i = 0; i < n; i++) {
        assert(1 <= m[i]);
        Int r1 = r[i] % m[i], m1 = m[i];
        if (r1 < 0) r1 += m1;
        if (m0 < m1) {
            std::swap(r0, r1);
            std::swap(m0, m1);
        }
        if (m0 % m1 == 0) {
            if (r0 % m1 != r1) return {0, 0};
            continue;
        }
        Int g, im;
        std::tie(g, im) = inv_gcd<Int>(m0, m1);

        Int u1 = m1 / g;
        if ((r1 - r0) % g) return {0, 0};

        Int x = (r1 - r0) / g % u1 * im % u1;
        r0 += x * m0;
        m0 *= u1;
        if (r0 < 0) r0 += m0;
    }
    return {r0, m0};
}

// 蟻本 P.262
// 中国剰余定理を利用して，色々な素数で割った余りから元の値を復元
// 連立線形合同式 A * x = B mod M の解
// Requirement: M[i] > 0
// Output: x = first MOD second (if solution exists), (0, 0) (otherwise)
template <class Int>
std::pair<Int, Int>
linear_congruence(const std::vector<Int> &A, const std::vector<Int> &B, const std::vector<Int> &M) {
    Int r = 0, m = 1;
    assert(A.size() == M.size());
    assert(B.size() == M.size());
    for (int i = 0; i < (int)A.size(); i++) {
        assert(M[i] > 0);
        const Int ai = A[i] % M[i];
        Int a = ai * m, b = B[i] - ai * r, d = std::__gcd(M[i], a);
        if (b % d != 0) {
            return std::make_pair(0, 0); // 解なし
        }
        Int t = b / d * mod_inverse<Int>(a / d, M[i] / d) % (M[i] / d);
        r += m * t;
        m *= M[i] / d;
    }
    return std::make_pair((r < 0 ? r + m : r), m);
}

template <class Int = int, class Long = long long> Int pow_mod(Int x, long long n, Int md) {
    static_assert(sizeof(Int) * 2 <= sizeof(Long), "Watch out for overflow");
    if (md == 1) return 0;
    Int ans = 1;
    while (n > 0) {
        if (n & 1) ans = (Long)ans * x % md;
        x = (Long)x * x % md;
        n >>= 1;
    }
    return ans;
}
#line 5 "number/combination.hpp"

// nCr mod m = p^q (p: prime, q >= 1)
// Can be used for n, r <= 1e18, m <= 1e7
// Complexity: O(m) (construction), O(log(n)) (per query)
// https://ferin-tech.hatenablog.com/entry/2018/01/17/010829
struct combination_prime_pow {
    int p, q, m;
    std::vector<int> fac, invfac, ppow;

    long long _ej(long long n) const {
        long long ret = 0;
        while (n) ret += n, n /= p;
        return ret;
    }

    combination_prime_pow(int p_, int q_) : p(p_), q(q_), m(1), ppow{1} {
        for (int t = 0; t < q; ++t) m *= p, ppow.push_back(m);
        fac.assign(m, 1);
        invfac.assign(m, 1);
        for (int i = 1; i < m; ++i) fac[i] = (long long)fac[i - 1] * (i % p ? i : 1) % m;
        invfac[m - 1] = fac[m - 1]; // Same as Wilson's theorem
        assert(1LL * fac.back() * invfac.back() % m == 1);
        for (int i = m - 1; i; --i) invfac[i - 1] = (long long)invfac[i] * (i % p ? i : 1) % m;
    }

    int nCr(long long n, long long r) const {
        if (r < 0 or n < r) return 0;
        if (p == 2 and q == 1) return !((~n) & r); // Lucas
        long long k = n - r;
        long long e0 = _ej(n / p) - _ej(r / p) - _ej(k / p);
        if (e0 >= q) return 0;

        long long ret = ppow[e0];
        if (q == 1) { // Lucas
            while (n) {
                ret = __int128(ret) * fac[n % p] * invfac[r % p] * invfac[k % p] % p;
                n /= p, r /= p, k /= p;
            }
            return (int)ret;
        } else {
            if ((p > 2 or q < 3) and (_ej(n / m) - _ej(r / m) - _ej(k / m)) & 1) ret = m - ret;
            while (n) {
                ret = __int128(ret) * fac[n % m] * invfac[r % m] * invfac[k % m] % m;
                n /= p, r /= p, k /= p;
            }
            return (int)ret;
        }
    }
};

// nCr mod m
// Complexity: O(m) space worst (construction), O(log(n) log(m)) (per query)
// Input: pairs of (prime, degree), such as vector<pair<int, int>> and map<int, int>
// https://judge.yosupo.jp/problem/binomial_coefficient
struct combination {
    std::vector<combination_prime_pow> cpps;
    std::vector<int> ms;

    template <class Map> combination(const Map &p2deg) {
        for (auto f : p2deg) {
            cpps.push_back(combination_prime_pow(f.first, f.second));
            ms.push_back(cpps.back().m);
        }
    }

    int operator()(long long n, long long r) const {
        if (r < 0 or n < r) return 0;
        std::vector<int> rs;
        for (const auto &cpp : cpps) rs.push_back(cpp.nCr(n, r));
        return crt(rs, ms).first;
    }
};

Back to top page