cp-library

競技プログラミング用 C++ ライブラリ

View the Project on GitHub qandori/cp-library

:heavy_check_mark: modint998244353 の階乗・逆階乗生成

components/mint_factorials.hpp · GitHub で見る

共通ページ

FPS と二項係数の前計算で共有する内部処理です。NTT には依存せず、共通の剰余演算を使います。通常は 階乗・逆階乗と二項係数 を利用します。

mint_detail::factorials(fact,invfact,n) は fact[i] = i!、invfact[i] = 1/i! を先頭 n 要素へ書き込みます。mint* 版では \(0\le n\le998244353\)、内部 Montgomery ワードのポインタ版では \(0\le n<998244353\) とし、それぞれ n 要素を持つ重ならない出力領域を渡します。n = 0 では書き込みません。

生成は \(O(n)\) 時間・\(O(1)\) 追加空間と逆元計算 1 回です。

依存ライブラリ

検証コード

コード

#ifndef MINT_FACTORIALS_HPP
#define MINT_FACTORIALS_HPP

#include "mint_arithmetic.hpp"

namespace mint_detail {

#if defined(__x86_64__) || defined(__i386__)

__attribute__((target("avx2"))) inline void factorial_transpose(Ops::I8 (&x)[8]) {
	Ops::I8 a[8], b[8];
	for (int i = 0; i < 8; i += 2) {
		a[i] = _mm256_unpacklo_epi32(x[i], x[i + 1]);
		a[i + 1] = _mm256_unpackhi_epi32(x[i], x[i + 1]);
	}
	for (int i = 0; i < 8; i += 4) {
		b[i] = _mm256_unpacklo_epi64(a[i], a[i + 2]);
		b[i + 1] = _mm256_unpackhi_epi64(a[i], a[i + 2]);
		b[i + 2] = _mm256_unpacklo_epi64(a[i + 1], a[i + 3]);
		b[i + 3] = _mm256_unpackhi_epi64(a[i + 1], a[i + 3]);
	}
	for (int i = 0; i < 4; ++i) {
		x[i] = _mm256_permute2x128_si256(b[i], b[i + 4], 0x20);
		x[i + 4] = _mm256_permute2x128_si256(b[i], b[i + 4], 0x31);
	}
}

template<bool reverse>
__attribute__((target("avx2"))) inline void factorial_chain(Ops::U* out, int chunk, Ops::I8 value, Ops::I8 factor) {
	using I = Ops::I8;
	I step = _mm256_set1_epi32(Ops::one);
	I order = _mm256_setr_epi32(7, 6, 5, 4, 3, 2, 1, 0);
	for (int j = 0; j < chunk; j += 8) {
		I block[8];
		for (int k = 0; k < 8; ++k) {
			block[k] = value;
			value = Ops::mul8(value, factor);
			if constexpr (reverse) factor = Ops::sub8(factor, step);
			else factor = Ops::add8(factor, step);
		}
		factorial_transpose(block);
		for (int k = 0; k < 8; ++k) {
			if constexpr (reverse) Ops::store8(out + (k + 1) * chunk - j - 8, _mm256_permutevar8x32_epi32(block[k], order));
			else Ops::store8(out + k * chunk + j, block[k]);
		}
	}
}

__attribute__((target("avx2"))) inline void factorials8(Ops::U* fact, Ops::U* invfact, int n) {
	using I = Ops::I8;
	using U = Ops::U;
	const int chunk = n / 256 * 8, covered = chunk * 32;
	alignas(32) U factor[32], products[32], starts[33], inverse[33];
	I p[4], x[4], step = _mm256_set1_epi32(Ops::one);
	for (int i = 0; i < 32; ++i) factor[i] = Ops::in(i * chunk + 1);
	for (int k = 0; k < 4; ++k) p[k] = _mm256_set1_epi32(Ops::one), x[k] = Ops::load8(factor + 8 * k);
	for (int i = 0; i < chunk; ++i) {
		for (int k = 0; k < 4; ++k) p[k] = Ops::mul8(p[k], x[k]), x[k] = Ops::add8(x[k], step);
	}
	for (int k = 0; k < 4; ++k) Ops::store8(products + 8 * k, p[k]);
	starts[0] = Ops::one;
	for (int i = 0; i < 32; ++i) starts[i + 1] = Ops::mul(starts[i], products[i]);
	U last = starts[32];
	for (int i = covered; i < n; ++i) fact[i] = last, last = Ops::mul(last, Ops::in(i + 1));
	last = mint::word(last).inv().v;
	for (int i = n - 1; i >= covered; --i) last = Ops::mul(last, Ops::in(i + 1)), invfact[i] = last;
	inverse[32] = last;
	for (int i = 31; i >= 0; --i) inverse[i] = Ops::mul(inverse[i + 1], products[i]);
	for (int k = 0; k < 4; ++k) {
		factorial_chain<false>(fact + 8 * k * chunk, chunk, Ops::load8(starts + 8 * k), Ops::load8(factor + 8 * k));
		alignas(32) U values[8], factors[8];
		for (int j = 0; j < 8; ++j) {
			const int end = (8 * k + j + 1) * chunk;
			values[j] = Ops::mul(inverse[8 * k + j + 1], Ops::in(end));
			factors[j] = Ops::in(end - 1);
		}
		factorial_chain<true>(invfact + 8 * k * chunk, chunk, Ops::load8(values), Ops::load8(factors));
	}
}
#endif

inline void factorials(Ops::U* fact, Ops::U* invfact, int n) {
	if (!n) return;
#if defined(__x86_64__) || defined(__i386__)
	if (n >= 256 && Ops::vectorized()) {
		factorials8(fact, invfact, n);
		return;
	}
#endif
	Ops::U value = Ops::one;
	for (int i = 0; i < n; ++i) fact[i] = value, value = Ops::mul(value, Ops::in(i + 1));
	value = mint::word(fact[n - 1]).inv().v;
	for (int i = n - 1; i >= 0; --i) invfact[i] = value, value = Ops::mul(value, Ops::in(i));
}

template<class T> requires std::is_same_v<T, mint>
inline void factorials(T* fact, T* invfact, int n) {
	assert(0 <= n && unsigned(n) <= mint::mod());
	static_assert(sizeof(mint) == sizeof(Ops::U) && std::is_trivially_copyable_v<mint>);
	if (unsigned(n) == mint::mod()) {
		factorials(fact, invfact, n - 1);
		fact[n - 1] = -fact[n - 2];
		invfact[n - 1] = -invfact[n - 2];
		return;
	}
	factorials(reinterpret_cast<Ops::U*>(fact), reinterpret_cast<Ops::U*>(invfact), n);
}

}

#endif
Back to top page