はじめに

この記事は、私のブログ「任意 mod で二項係数を列挙する (2)」の内容を整理したものになります。

次節で示す問題を $\displaystyle O\left(\dfrac{N\log M}{\log \log M}\right)$ で解くアルゴリズムを解説します。

問題

整数 $N,M$ が与えられる。全ての $k=0,\ldots,N$ に対して、二項係数
$$
{}_N C _{k} =\frac{N!}{k!(N-k)!}
$$
の値を $\mathrm{mod}\; M$ で求めてください。

制約

$0\leq N\leq 10^7$
$1\leq M\leq 10^9$

準備

二項係数 ${}_N C _{k}$ に関して式 $(1)$ が成立するので、この漸化式に従って $k=0,\ldots,N$ の順に求めていくことを考えます。

$$
{}_N C _{k}=\begin{cases}
1& \text{if $\;k=0$ }\\
\dfrac{N-k+1}{k}\cdot {}_N C _{k-1} & \text{otherwise}
\end{cases}\tag{1}
$$

$M$ が $N$ よりも大きな素数であれば、任意の $k=1,\ldots,N$ に対して $\mathrm{mod}\; M$ における乗法逆元が存在します。従って、この場合は「1000000007 で割ったあまり」の求め方を総特集！〜逆元から離散対数まで〜で紹介されている方法により容易に問題を解くことが出来ます。

ここで、$\mathrm{mod}\; M$ における整数 $a$ の乗法逆元とは、式 $(2)$ 満たす整数 $x$ を指します。

$$
a\cdot x\equiv 1 \pmod M \tag{2}
$$

本記事を通して法 $M$ は固定なので、以降はこの $x$ を単に「$a^{-1}$」や「$a$ の逆元」と表記することがあります。

今回考えている問題では $M\leq N$ であるケースや $M$ が合成数のケースも存在するため、任意の $k=1,\ldots,N$ に対して $k$ の逆元が存在するとは限りません。一方で、必ず存在しないというわけでもありません (例: $3\cdot 3\equiv 1\pmod 4$)。ではどのような $k$ に対して逆元が存在するかというと、次の事実が成立します。詳しい証明については一次不定方程式ax+by=cの整数解に譲ります。

$$
\text{$k$ の $\mathrm{mod}\; M$ における乗法逆元が存在} \iff \gcd(k,M)=1.
$$

結局、この節では「$k$ が $M$ と共通の素因数を持つ場合は $\mathrm{mod}\; M$ におけるモジュラ逆数 $k^{-1}$ が存在しないのでマズイ」ということを言いたかったと理解してください。

解法のアイデア

上述の通り、$k$ と $M$ が素因数を共有する場合は漸化式 $(1)$ に従った素直な更新が難しいということでした。そこで、以下では $k$ を次のように「分解」して考えます。

$$
k=(\text{$M$ と素因数を共有しない部分})\times(\text{$M$ と素因数を共有する部分})
$$

以降、$M$ と素因数を共有しない部分を $S_M(k)$、$M$ と素因数を共有する部分を $T_M(k)$ とします。

例えば $k=360=2^3 \cdot 3^2 \cdot 5,\ M=70=2\cdot 5\cdot 7$ のとき、$k$ と $M$ に共通する素因数は $2,5$ なので、$k$ は次のように分解されます。
$$
k=\color{red}{\underline{\color{black}{9}}}\times \color{blue}{\underline{\color{black}{2^3\cdot 5}}}
$$
ここで、赤線を引いた部分が $M$ と素因数を共有しない部分であり、青線を引いた部分が $M$ と素因数を共有する部分です。つまり、この例では $S_M(k)=9,\ T_M(k)=2^3\cdot 5$ です。

さて、このようにして $k=S_M(k)\times T_M(k)$ と分解したとき、定義から明らかに $S_M(k)$ に関しては逆元 $S_M(k)^{-1}$ が存在します。逆元の具体的な計算方法は一旦置いておくことにすれば、あとは $T_M(k)$ の部分をどうにかしたいということになります。

$T_M(k)$ を考えるにあたって、まずは $M$ が持つ素因数を知る必要があります。そこで、$M$ が $L$ 個の素数 $p_1\lt p_2\lt \cdots\lt p_L$ を用いて次のように素因数分解されるとします。
$$
M=\prod_{i=1}^{L}p_i^{q_i}\quad (\text{$q_i$ は正整数})
$$
正整数 $A$ および $i\in\{1,\ldots,L\}$ に対して、$q_i(A)$ を $A$ の素因数分解における $p_i$ の指数として定めます。$A$ が素因数 $p_i$ を持たない場合は $q_i(A)=0$ とします。ここで、$T_M$ の定義から、式 $(3)$ が成立します。
$$
T_M(A)=\prod_{i=1}^{L}p_i^{q_i(A)}\tag{3}
$$

本手法では、$T_M({}_N C _k)$ の代わりに $q_i({}_N C _k)\;(i=1,\ldots,L)$ を管理します。つまり、計算中の ${}_N C _k$ の値を以下のような形で保持します。
$$
{} _ N C _ k = S _ M( {} _ N C _ k ) \times \prod _ {i=1}^{L}p_i^{q_i({}_N C _ k )}
$$

ところで、$M$ の素因数分解の時点で冒頭の計算量を達成できなくなるような気がしますが、実は問題ありません。説明の都合上 $M$ の完全な素因数分解形を用いましたが、${}_NC_k$ は $N$ 以下の素因数しか持たないので、$M$ が持つ $N$ 以下の素因数さえ求まっていれば正しく分解することが出来ます。

解法

上記のアイデアを具体的な実装に落としていきます。まず、線形篩 ¹ により $N$ 以下の素数を $O(N)$ で求めておきます。そして、得られた素数のリストを用いて $M$ の $N$ 以下の素因数 $p_1\lt p_2\lt\cdots\lt p_L$ を計算します。このとき、同時に各 $v=1,\ldots,N$ に対して次で定義される $d(v)$ を計算します。
$$
d(v)=\max\{i\mid \text{$p_i$ は $v$ を割り切る} \}\cup\{0\}
$$
これは、擬似的には次のように求められます。

// N 以下の素数を昇順に並べた列 primes は線形篩により計算済みとする
// 10^9 以下の整数が持つ素因数の種類数の最大値
int MAX_PRIME_FACTOR_NUM = 9; 
int[] d = new int[N + 1];
int L = 0;
for (int v = 0; v <= N; ++v) d[v] = L;
int[] p = new int[MAX_PRIME_FACTOR_NUM + 1];
for (int prime : primes) {
    // M を割り切る素数を発見
    if (M % prime == 0) {
        ++L;
        p[L] = prime;
        // prime の倍数であるような v に対して d[v] = L で更新
        for (int v = prime; v <= N; v += prime) {
            d[v] = L;
        }
    }
}

計算量を解析します。まず，$M$ が持つ素因数の種類数 $L$ に関して、$L=O\left(\dfrac{\log M}{\log \log M}\right)$ が成り立ちます²。そして、$L$ 個の素数 $p_1,\ldots,p_L$ に対して内側のループは $\lfloor N/p_i\rfloor$ 回だけ回ります。従って、計算量は次のように評価されます。
$$
\sum_{i=1}^L\left\lfloor\frac{N}{p_i}\right\rfloor\leq\sum_{i=1}^L N=NL=O\left(\frac{N\log M}{\log \log M}\right)
$$
なお、擬似コード中の MAX_PRIME_FACTOR_NUM = 9 は $2\times 3\times\ 5\times 7\times 11\times 13\times 17\times 19 \times 23 \times 29\simeq 6.47\cdot 10^9$ に基づいています。

求めた $d$ を用いることで、$k$ に対して次のように $T_M(k)$ を計算出来ます。$\displaystyle T_M(k)=\prod_{i=1}^{L}p_i^{q_i(k)}$ の形で求めたいので、配列として表現していることに注意です。

int[] t = new int[MAX_PRIME_FACTOR_NUM + 1];
for (int i = 0; i <= MAX_PRIME_FACTOR_NUM; ++i) {
    t[i] = 0;
}
int cur = k;
while (true) {
    // d の定義から、M の素因数 p_1, ..., p_L のうち、cur を割り切る最大の素数が p_i
    int i = d[cur];
    // i = 0 の場合は、そのような p_i が存在しなかったことを表す。
    if (i == 0) break;
    ++t[i];
    cur /= p[i];
}

式 $(1)$ に従った更新をするためには、$k=1,\ldots,N$ に対して $T_M(k)$ および $T_M(N-k+1)$ を計算する必要があります。つまり、すべての $k=1,\ldots,N$ に対して $2$ 回ずつ上記の計算が行われます。$2$ 回ずつというのは定数倍なので無視することにして、$k=1,\ldots,N$ に対して $1$ 回ずつ計算する場合の計算量を解析します。

ある $j\in\{1,\ldots,L\}$ を固定します。while の中で $i=j$ となるのは $\displaystyle \sum_{k=1}^N q_j(k)$ 回であり、これは $N!$ が $p_j$ で割り切れる回数に一致します。この回数は $\displaystyle\sum_{n=1}^{\infty}\left\lfloor\frac{N}{p_j^n}\right\rfloor$ であることが知られており ³、次のように評価することが出来ます。
$$
\sum_{n=1}^{\infty}\left\lfloor\frac{N}{p_j^n}\right\rfloor\leq
\sum_{n=1}^{\infty}\frac{N}{p_j^n}=\frac{1}{p_j}\cdot\frac{N}{1-1/p_j}=\frac{N}{p_j-1}\leq N\quad (\because p_j\geq 2)
$$
両辺 $j=1,\ldots,L$ について和を取ることで、次が成り立ちます。
$$
\sum_{j=1}^L\sum_{n=1}^{\infty}\left\lfloor\frac{N}{p_j^n}\right\rfloor\leq NL=O\left(\frac{N\log M}{\log \log M}\right)
$$
上記のコードにおいて、副産物として while ループを抜けた後は cur の値が $S_M(k)$ と等しくなっています。

以上で各 $k=1,\ldots,N$ に対して $S_M(k)$ および $T_M(k)$ を $\displaystyle O\left(\frac{N\log M}{\log \log M}\right)$ で求めることが出来ました。以下のコードでは、$S_M(k)$ を s[k]、$T_M(k)$ を t[k] と表します。t は各 $k$ に対して $q_1(k),\ldots,q_L(k)$ を保持しているので $2$ 次元配列であることに注意してください。

続いて、$S_M({}_NC _{k})$ および $T_M({}_NC _{k})$ を計算していきます。以下のコードでは $S_M({}_NC _{k})$ を S、$T_M({}_NC _{k})$ を T と表します。くどいようですが、T は配列です。

式 $(1)$ に従って、次のように求めていくことが出来ます。

$k=0$ のとき、S=1、T[i]=0 for i=1,...,L で初期化する。
$k\gt 0$ のとき、式 $(1)$ に従い S*=s[n-k+1]/s[k]、T[i]+=t[n-k+1][i]-t[k][i] for i=1,...,L と更新する。ただし、S の更新において除算 / は $\mathrm{mod}\; M$ における除算であることに注意する。つまり、s[k] の逆元を求めておく必要がある。

上の通り、$S_M(k)$ の逆元が必要となります。$S_M(k)\leq N$ なので、$N$ 以下の整数に対して $\mathrm{mod}\; M$ での逆元を予め求めておくことにしましょう。これは、脚注 [1] の記事で紹介されているアルゴリズムにより $O(N)$ で可能です。以上より、s や t が計算済みという仮定の下で、S および T の更新は全体 $\displaystyle O\left(\frac{N\log M}{\log \log M}\right)$ で可能です。

$S_M({}_NC _k)$ と $T_M({}_NC _k)$ は求まりましたが、最後の仕事として ${}_NC _k$ を復元する必要があります。定義を思い出すと、
$$
{}_NC _k=S_M({}_NC _k)\times\prod _{i=1}^{L}p_i^{q_i({}_NC _k)}
$$
でした。総積記号 ($\prod$) の中の累乗の計算を各々定数時間で済ますことができれば、全体として $\displaystyle O\left(\frac{N\log M}{\log \log M}\right)$ が達成されます。そして、これは可能です。

計算に必要な累乗を全て前計算することを考えます。まず、累乗の底は $p_1,\ldots,p_L$ に限られています。また、${}_NC _k$ に含まれる素因数 $p_i$ の個数 $q_i({}_NC _k)$ は、${}_NC _k$ が $N!$ を割り切ることから以下のように評価されます。
$$
q_i({}_NC _k)\leq q_i(N!)=\sum _{n=1}^{\infty}\left\lfloor\frac{N}{p_i^n}\right\rfloor\leq\sum _{n=1}^{\infty}\frac{N}{p_i^n}=\frac{N}{p_i-1}
$$
つまり、$p_i$ に対して $p_i^0,\ldots,p_i^{\lfloor N/(p_i-1)\rfloor}$ を前計算しておけば十分です。そして、
$$
\sum _{i=1}^L 1+\left\lfloor \frac{N}{p_i-1}\right\rfloor\leq \sum _{i=1}^L (1+N)=O\left(\frac{N\log M}{\log \log M}\right) \quad(\because p_i\geq 2)
$$
より、前計算にかかる計算量はは $\displaystyle O\left(\frac{N\log M}{\log \log M}\right)$ です。

以上より、全ての計算を $\displaystyle O\left(\frac{N\log M}{\log \log M}\right)$ で行うことが出来ました。なお、実用上は ${}_NC _k={}_NC _{N-k}$ を用いて後ろ半分の計算をサボることで定数倍高速化を図ることができます。

実装例 (C++)

verified: AtCoder Regular Contest 012 D - Don't worry. Be Together

折り畳み

#include <vector>
#include <numeric>

// referece: https://37zigen.com/linear-sieve/
class LinearSieve {
    public:
        LinearSieve(unsigned int n) : _n(n), min_prime_factor(std::vector<unsigned int>(n + 1)) {
            std::iota(min_prime_factor.begin(), min_prime_factor.end(), 0);
            prime_list.reserve(_n + 1);
            for (unsigned int d = 2; d <= _n; ++d) {
                if (min_prime_factor[d] == d) {
                    prime_list.push_back(d);
                }
                unsigned int prime_max = std::min(min_prime_factor[d], _n / d);
                for (unsigned int prime : prime_list) {
                    if (prime > prime_max) {
                        break;
                    }
                    min_prime_factor[prime * d] = prime;
                }
            }
        }
        unsigned int prime_num() const {
            return prime_list.size();
        }
        const std::vector<unsigned int>& get_prime_list() const {
            return prime_list;
        }
        const std::vector<unsigned int>& get_min_prime_factor() const {
            return min_prime_factor;
        }
    private:
        const unsigned int _n;
        std::vector<unsigned int> min_prime_factor;
        std::vector<unsigned int> prime_list;
};

// nCk mod m を全ての k = 0, ..., n に対して求める．m は素数でなくてもよい．
class ArbitraryModBinomialCoefficients {
    public:
        // nCk mod m を求める
        ArbitraryModBinomialCoefficients(unsigned int N, unsigned int M) :
            _N(N), _M(M), _sieve(LinearSieve(N)), _binom(std::vector<long long>(N + 1)) {
            solve();
        }
        // return nCk
        long long operator[](unsigned int k) const {
            return _binom[k];
        }
        // return nC0, nC1, ..., nCn mod m を格納した vector
        const std::vector<long long>& get_coeffs() const {
            return _binom;
        }
        // 副産物としての線形篩
        const LinearSieve& get_sieve() const {
            return _sieve;
        }
    private:
        const unsigned int _N, _M;
        const LinearSieve _sieve;
        std::vector<long long> _binom;

        // return y = (x mod M) (0 <= y < M)
        inline unsigned int safe_mod(const long long x) {
            int y = x % _M;
            return y < 0 ? y + _M : y;
        }

        // return x s.t. 0 <= x < m and (v * x mod m) = gcd(a, m)
        unsigned int gcd_inv(const unsigned int v) {
            long long a = v, b = _M;
            long long x = 1, z = 0;
            long long y = 0, w = 1;
            long long tmp;
            while (b) {
                long long p = a / b, q = a % b;
                tmp = x - z * p; x = z; z = tmp;
                tmp = y - w * p; y = w; w = tmp;
                a = b; b = q;
            }
            return safe_mod(x);
        }

        // reference: https://37zigen.com/linear-sieve/
        // 合成数 mod に対して逆元のテーブルを求める
        // 逆元が存在しない添え字に対して，値は未定義
        void mod_invs(std::vector<long long>& invs) {
            auto &mpf = _sieve.get_min_prime_factor();
            for (unsigned int i = 0; i <= _N; ++i) {
                unsigned int pf = mpf[i];
                if (pf == i) {
                    invs[i] = gcd_inv(i);
                } else {
                    invs[i] = invs[pf] * invs[i / pf] % _M;
                }
            }
        }

        void solve() {
            // sieve は線形篩
            // 最小素因数のテーブルと素数のリストを持つ
            auto &primes = _sieve.get_prime_list();
            // d と p を求める
            std::vector<unsigned int> d(_N + 1, 0);
            std::vector<unsigned int> p;
            for (unsigned int prime : primes) {
                if (_M % prime) continue;
                p.push_back(prime);
                unsigned int sz = p.size();
                for (unsigned int v = prime; v <= _N; v += prime) {
                    d[v] = sz;
                }
            }
            const unsigned int L = p.size();
            // p は 1-indexed なのでダミー要素を先頭に入れる
            p.insert(p.begin(), 0);
            // invs[k] = k ^ {-1} mod M (存在する場合)
            // 存在しない場合の値は未定義だが，アクセスすることはない．
            std::vector<long long> invs(_N + 1);
            mod_invs(invs);
            // べき乗の前計算
            std::vector<std::vector<long long>> powers(L + 1);
            for (unsigned int i = 1; i <= L; ++i) {
                // 指数 (べき乗の肩の値) の上界
                unsigned int max_index = _N / (p[i] - 1);
                powers[i].resize(max_index + 1);
                powers[i][0] = 1;
                for (unsigned int j = 0; j < max_index; ++j) {
                    powers[i][j + 1] = powers[i][j] * p[i] % _M;
                }
            }
            // 定数倍高速化のため，前半分だけを漸化式に従って計算する
            const unsigned int half = (_N + 1) / 2;
            // k = 0 の場合について，S, T を初期化
            long long S = 1;
            std::vector<unsigned int> T(L + 1, 0);
            // nC0 = 1
            _binom[0] = 1;
            for (unsigned int k = 1; k <= half; ++k) {
                // num: 漸化式の分子, den: 漸化式の分母
                unsigned int num = _N - k + 1, den = k;
                // T の更新
                // t を求めず，直接 T に足しこむことで余分な領域を使わずに済ませる
                while (d[num]) ++T[d[num]], num /= p[d[num]];
                while (d[den]) --T[d[den]], den /= p[d[den]];
                // S の更新
                S = S * num % _M;
                S = S * invs[den] % _M;
                // S, T から nCk を復元
                _binom[k] = S;
                for (unsigned int i = 1; i <= L; ++i) {
                    _binom[k] = _binom[k] * powers[i][T[i]] % _M;
                }
            }
            // 後ろ半分は nCk = nCn-k を用いて求めると定数倍がいい
            for (unsigned int k = half + 1; k <= _N; ++k) {
                _binom[k] = _binom[_N - k];
            }
        }
};

AtCoder Library を用いることが出来る場合は、mint に atcoder::modint を渡すことを想定して次のように実装を簡単化することができます。

verifid : AtCoder Regular Contest 012 D - Don't worry. Be Together

折り畳み

#include <vector>
#include <numeric>

// 線形篩の部分は同じなので省略

template <typename mint>
class ArbitraryModBinomialCoefficients {
    public:
        ArbitraryModBinomialCoefficients(unsigned int N) :
            _N(N), _M(mint::mod()), _sieve(LinearSieve(N)), _binom(std::vector<mint>(N + 1)) {
            solve();
        }
        mint operator[](unsigned int k) const {
            return _binom[k];
        }
        const std::vector<mint>& get_coeffs() const {
            return _binom;
        }
        const LinearSieve& get_sieve() const {
            return _sieve;
        }
    private:
        const unsigned int _N, _M;
        const LinearSieve _sieve;
        std::vector<mint> _binom;

        void mod_invs(std::vector<mint>& invs) {
            auto &mpf = _sieve.get_min_prime_factor();
            if (_N >= 1) invs[1] = 1;
            for (unsigned int i = 2; i <= _N; ++i) {
                unsigned int pf = mpf[i];
                if (pf == i) {
                    if (_M % pf) invs[i] = mint(i).inv();
                } else {
                    invs[i] = invs[pf] * invs[i / pf];
                }
            }
        }

        void solve() {
            auto &primes = _sieve.get_prime_list();
            std::vector<unsigned int> d(_N + 1, 0);
            std::vector<unsigned int> p;
            for (unsigned int prime : primes) {
                if (_M % prime) continue;
                p.push_back(prime);
                unsigned int sz = p.size();
                for (unsigned int v = prime; v <= _N; v += prime) {
                    d[v] = sz;
                }
            }
            const unsigned int L = p.size();
            p.insert(p.begin(), 0);
            std::vector<mint> invs(_N + 1);
            mod_invs(invs);
            std::vector<std::vector<mint>> powers(L + 1);
            for (unsigned int i = 1; i <= L; ++i) {
                unsigned int max_index = _N / (p[i] - 1);
                powers[i].resize(max_index + 1);
                mint pi = p[i];
                powers[i][0] = 1;
                for (unsigned int j = 0; j < max_index; ++j) {
                    powers[i][j + 1] = powers[i][j] * pi;
                }
            }
            const unsigned int half = (_N + 1) / 2;
            mint S = 1;
            std::vector<unsigned int> T(L + 1, 0);
            _binom[0] = 1;
            for (unsigned int k = 1; k <= half; ++k) {
                unsigned int num = _N - k + 1, den = k;
                while (d[num]) ++T[d[num]], num /= p[d[num]];
                while (d[den]) --T[d[den]], den /= p[d[den]];
                S *= num * invs[den];
                _binom[k] = S;
                for (unsigned int i = 1; i <= L; ++i) {
                    _binom[k] *= powers[i][T[i]];
                }
            }
            for (unsigned int k = half + 1; k <= _N; ++k) {
                _binom[k] = _binom[_N - k];
            }
        }
};

参考: https://37zigen.com/linear-sieve/ ↩
参考: http://matwbn.icm.edu.pl/ksiazki/aa/aa42/aa4242.pdf ↩
参考: https://manabitimes.jp/math/590 ↩