cp-library

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

View the Project on GitHub qandori/cp-library

:heavy_check_mark: Bostan–Mori の NTT 倍長化

components/fps_bostan.hpp · GitHub で見る

共通ページ

Linux x86_64 / GCC の AVX2 経路で使う内部部品です。通常は公開関数 Bostan–Mori または 線形漸化式 を使います。

操作 内容 時間計算量
fps_detail::bostan_detail::inverse_twisted(plan,in,out,powers,n) 正規化前の逆 NTT を in から out へ求め、各要素へ Montgomery 表現の powers[i] を掛ける \(O(n\log n)\)

n は8以上 \(2^{22}\) 以下の2冪とし、plan はその変換に必要な根表を持つものとします。in・out・powers は互いに重ならない各 n 要素の配列で、32バイト境界に整列させます。各要素は NTT::U の生存中のオブジェクトとし、内部ワードの範囲は \([0,2p)\)、\(p=998244353\) です。入力の変更や配列外の padding は要求しません。呼出側で NTT::vectorized() を確認します。

計算量: \(O(n\log n)\) 時間。入力・出力・係数倍・根表を含む空間は \(O(n)\)、局所変数は \(O(1)\) です。

依存ライブラリ

検証コード

コード

#ifndef CP_COMPONENTS_FPS_BOSTAN_HPP
#define CP_COMPONENTS_FPS_BOSTAN_HPP

#include "ntt.hpp"

#if defined(__linux__) && defined(__x86_64__) && defined(__GNUC__) && !defined(__clang__)
#pragma GCC push_options
_Pragma("GCC target(\"avx2\")")
#pragma GCC optimize("O3")

namespace fps_detail::bostan_detail {
inline void inverse_twisted(const NTT& plan, const NTT::U* in, NTT::U* out, const NTT::U* powers, int n) {
	using N = NTT;
	if (n == 8) {
		auto z = plan.small8<4, true>(plan.small8<2, true>(plan.small8<1, true>(N::load8(in), 0), 0), 0);
		N::store8(out, N::mul8(z, N::load8(powers)));
		return;
	}
	for (int j = 0; j < n; j += 8) {
		auto z = N::load8(in + j);
		z = plan.small8<4, true>(plan.small8<2, true>(plan.small8<1, true>(z, j / 8), j / 8), j / 8);
		N::store8(out + j, z);
	}
	int h = 8;
	for (; 4 * h < n; h *= 4) plan.radix4_8<true>(out, h, n);
	if (4 * h == n) {
		const auto im = _mm256_set1_epi32(plan.inverse_roots[1]);
		const auto iq = _mm256_set1_epi32(plan.inverse_root_quotients[1]);
		for (int i = 0; i < h; i += 8) {
			const auto x0 = N::load8(out + i), x1 = N::load8(out + h + i);
			const auto x2 = N::load8(out + 2 * h + i), x3 = N::load8(out + 3 * h + i);
			const auto u = N::add8(x0, x1), v = N::sub8(x0, x1), s = N::add8(x2, x3);
			const auto t = N::fixed_mul8(N::sub8(x2, x3), im, iq);
			N::store8(out + i, N::mul8(N::add8(u, s), N::load8(powers + i)));
			N::store8(out + h + i, N::mul8(N::add8(v, t), N::load8(powers + h + i)));
			N::store8(out + 2 * h + i, N::mul8(N::sub8(u, s), N::load8(powers + 2 * h + i)));
			N::store8(out + 3 * h + i, N::mul8(N::sub8(v, t), N::load8(powers + 3 * h + i)));
		}
	} else {
		for (int i = 0; i < h; i += 8) {
			const auto x = N::load8(out + i), y = N::load8(out + h + i);
			N::store8(out + i, N::mul8(N::add8(x, y), N::load8(powers + i)));
			N::store8(out + h + i, N::mul8(N::sub8(x, y), N::load8(powers + h + i)));
		}
	}
}
}

#pragma GCC pop_options
#endif
#endif
Back to top page