任意 mod の二項係数

math/binomial_mod.hpp · GitHub で見る

共通ページ

素数・素数冪・合成数を法として二項係数を求めます。法が 1 の場合も扱えます。返り値は modint ではなく、法で割った余りの整数です。

#include "lib/math/binomial_mod.hpp"

int main() {
    BinomialMod composite(72);
    auto a = composite.comb(10, 4); // 210 mod 72 = 66
    auto b = composite(1000000000000000000LL, 2); // 0
    auto c = composite(0, 0); // 1
    auto d = composite(3, 5); // 0

    BinomialMod prime(1000000007, 200000);
    auto e = prime(200000, 1); // 200000
}

操作	内容・返り値
BinomialMod(modulus,table_limit=1000000)	std::uint32_t の法と、素数冪ごとの前計算上限を指定して構築
mod()	法を std::uint32_t で返す
max_n()	問い合わせ可能な n の上限を std::int64_t で返す
comb(n,k) / operator()(n,k)	std::int64_t の n,k に対し、二項係数の余りを std::uint32_t で返す

法は 1 <= modulus <= UINT32_MAX とします。modulus==0 では std::invalid_argument を送出します。返り値は [0,modulus) に入り、法が 1 なら常に 0 です。n<0、k<0、n<k の場合も 0 を返します。

有効な問い合わせでは 0 <= n <= max_n() とします。範囲を超えると std::out_of_range を送出します。範囲外の k に対する 0 の返却と、法が 1 の処理を先に行うため、これらは n の上限にかかわらず使えます。

table_limit と n の上限について

法の素因数分解に現れる各素数冪を \(q=p^e\) とすると、それぞれ min(table_limit,q-1) まで前計算します。

• 全ての素数冪で q-1 <= table_limit なら、max_n() は INT64_MAX です。既定値では、法が 1000000 以下なら必ずこの条件を満たし、n=10^18 のような大きな値を扱えます。
• それ以外では、max_n() は table_limit です。例えば BinomialMod(1000000007,200000) は n<=200000 の問い合わせを扱います。
• 法が 1 の場合は前計算を行わず、max_n() は INT64_MAX です。

大きな素数の法では、必要な n に合わせて table_limit を指定してください。table_limit=0 も指定できます。前計算量は素数冪ごとの上限に比例するため、大きな上限には相応の時間・メモリが必要です。

#include "lib/math/binomial_mod.hpp"
#include <limits>

int main() {
    BinomialMod small_mod(1000);
    bool large_n_supported = small_mod.max_n() == std::numeric_limits<std::int64_t>::max(); // true
    BinomialMod bounded(1000000007, 500000);
    auto limit = bounded.max_n(); // 500000
}

各オブジェクトは固定した法・前計算表を持ち、異なる法のオブジェクトを同時に使えます。コピーは独立します。法や上限を変更する場合は、新しいオブジェクトを構築してください。既存の 階乗・逆階乗による二項係数 とは独立しており、このヘッダーの include だけで全体の階乗表を自動生成することはありません。

法を \(m\)、相異なる素因数の個数を \(r\)、各素数冪を \(q_i=p_i^{e_i}\)、前計算上限を \(L\)、\(S=\displaystyle\sum_{i=1}^{r}\min(L,q_i-1)\) とします。構築は \(O(\sqrt m+S+r\log m)\) 時間、保持空間は \(O(S+\log m)\) です。問い合わせは \(O(\displaystyle\sum_{i=1}^{r}\log_{p_i}n)\) 時間、作業空間は固定長の配列です。整数幅を固定として乗算・除算を定数時間とします。法が 1、範囲外の k、k=0、k=n は定数時間です。

依存ライブラリ

• 32 bit Barrett 演算
• 一次合同式

検証コード

• test/binomial_mod.test.cpp

コード

#ifndef BINOMIAL_MOD_HPP
#define BINOMIAL_MOD_HPP

#include "../components/barrett32.hpp"
#include "linear_congruence.hpp"
#include <algorithm>
#include <array>
#include <bit>
#include <limits>
#include <vector>

class BinomialMod {
	using U = std::uint32_t;
	using W = std::uint64_t;
	struct PrimePower {
		U prime, exponent, coefficient;
		Barrett32 modulus;
		std::vector<U> factorial, inverse, powers;
		PrimePower(U p, U e, U q, U limit, U weight)
			: prime(p), exponent(e), coefficient(weight), modulus(q),
			  factorial(std::size_t(std::min(limit, q - 1)) + 1), inverse(factorial.size()), powers(e, 1) {
			factorial[0] = 1;
			for (std::size_t i = 1; i < factorial.size(); ++i)
				factorial[i] = i % p ? modulus.mul(factorial[i - 1], U(i)) : factorial[i - 1];
			inverse.back() = U(linear_congruence(factorial.back(), 1, q)->residue);
			for (std::size_t i = factorial.size() - 1; i; --i)
				inverse[i - 1] = i % p ? modulus.mul(inverse[i], U(i)) : inverse[i];
			for (U i = 1; i < e; ++i) powers[i] = U(W(powers[i - 1]) * p);
		}
		template<bool Binary> U choose_power(W n, W k) const {
			struct Level { W n, k, r; };
			std::array<Level, 64> levels;
			W r = n - k;
			U valuation = 0;
			bool negative = false;
			if constexpr (Binary) {
				valuation = U(std::popcount(k) + std::popcount(r) - std::popcount(n));
				if (valuation >= exponent) return 0;
				negative = exponent == 2 && ((valuation - ((n >> 1) - (k >> 1) - (r >> 1))) & 1);
			}
			U length = 0;
			while (n) {
				levels[length++] = {n, k, r};
				if constexpr (Binary) n >>= 1, k >>= 1, r >>= 1;
				else {
					n /= prime, k /= prime, r /= prime;
					U carry = U(n - k - r);
					valuation += carry;
					if (valuation >= exponent) return 0;
					if (length >= exponent) negative ^= carry != 0;
				}
			}
			U a = 1, b = 1, c = 1;
			for (U i = 0; i < length; ++i) {
				auto value = levels[i];
				if constexpr (Binary) value.n &= modulus.mod - 1, value.k &= modulus.mod - 1, value.r &= modulus.mod - 1;
				else if (i + exponent < length) {
					auto q = levels[i + exponent];
					value.n -= q.n * modulus.mod;
					value.k -= q.k * modulus.mod;
					value.r -= q.r * modulus.mod;
				}
				a = modulus.mul(a, factorial[std::size_t(value.n)]);
				b = modulus.mul(b, inverse[std::size_t(value.k)]);
				c = modulus.mul(c, inverse[std::size_t(value.r)]);
			}
			U result = modulus.mul(modulus.mul(a, b), c);
			if (negative) result = modulus.mod - result;
			return modulus.mul(result, powers[valuation]);
		}
		U choose(W n, W k) const {
			if (exponent == 1) {
				if (prime == 2) return (n & k) == k;
				U result = 1;
				while (n) {
					U a = U(n % prime), b = U(k % prime);
					if (a < b) return 0;
					result = modulus.mul(result, modulus.mul(factorial[a], modulus.mul(inverse[b], inverse[a - b])));
					n /= prime, k /= prime;
				}
				return result;
			}
			return prime == 2 ? choose_power<true>(n, k) : choose_power<false>(n, k);
		}
	};
	Barrett32 modulus_;
	std::vector<PrimePower> factors_;
	std::int64_t max_n_ = std::numeric_limits<std::int64_t>::max();
	static U checked_(U m) {
		if (!m) throw std::invalid_argument("BinomialMod modulus");
		return m;
	}

public:
	explicit BinomialMod(U modulus, U table_limit = 1000000) : modulus_(checked_(modulus)) {
		auto add = [&](U p, U e, U q) {
			U weight = U(W(modulus / q) * W(linear_congruence(modulus / q, 1, q)->residue));
			factors_.emplace_back(p, e, q, table_limit, weight);
			if (q - 1 > table_limit) max_n_ = table_limit;
		};
		U remaining = modulus;
		for (U p = 2; p <= remaining / p; p += p == 2 ? 1 : 2) {
			if (remaining % p) continue;
			U q = 1, e = 0;
			do { remaining /= p; q *= p; ++e; } while (remaining % p == 0);
			add(p, e, q);
		}
		if (remaining > 1) add(remaining, 1, remaining);
	}
	U mod() const { return modulus_.mod; }
	std::int64_t max_n() const { return max_n_; }
	U comb(std::int64_t n, std::int64_t k) const {
		if (n < 0 || k < 0 || n < k || mod() == 1) return 0;
		if (n > max_n_) throw std::out_of_range("BinomialMod n");
		if (!k || n == k) return 1;
		if (factors_.size() == 1) return factors_[0].choose(W(n), W(k));
		U answer = 0;
		for (const auto &f : factors_)
			answer = modulus_.reduce(W(answer) + modulus_.mul(f.coefficient, f.choose(W(n), W(k))));
		return answer;
	}
	U operator()(std::int64_t n, std::int64_t k) const { return comb(n, k); }
};

#endif

Back to top page