cp-library

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

View the Project on GitHub qandori/cp-library

:heavy_check_mark: FPS の一点評価

components/fps_horner.hpp · GitHub で見る

共通ページ

fps_detail::horner(a,n,x) は、Mint の係数配列で表された多項式を一点 x で評価する内部部品です。通常は公開関数 fps::eval(x) を使います。係数は定数項から並べます。空の多項式は 0 を返し、x=0 のときは定数項を返します。入力は変更せず、追加の整列や padding を要求しません。

操作 内容・返り値 計算量
fps_detail::horner(a,n,x) a[0] から a[n-1] の多項式を x で評価 \(O(n)\)

通常は FPS の f.eval(x) を使います。例えば fps{1,2,3}.eval(2).val() は 17 です。

32 項以上では、指数を 8 で割った余りごとに Horner 法を分けます。AVX2 が使える場合は各組を SIMD の 8 lane で求め、64 項以上ではさらに 4 本の独立した SIMD の漸化式に分けます。係数との積には固定した x^8 または x^32 に対する Shoup 乗算を使い、最後に組ごとの値を Horner 法で結合します。共通の剰余演算部品で AVX2 の利用可否を実行時に判定し、NTT には依存しません。

Mint の内部ワードは \([0,2p)\)、\(p=998244353\) とします。係数と漸化式の値は Montgomery 表現のまま保持し、固定乗数だけを通常の剰余値にします。配列の末尾の不足分は所有する局所配列で 0 埋めします。係数の読み込みと SIMD 値の取り出しはオブジェクト表現のコピーで行い、配列外を読みません。

AVX2 が使えない環境の 32 項以上は、8 本の通常の Horner 法へ分けます。短い場合は直列の Horner 法を使います。FormalPowerSeries<T> の係数が内部の Mint と異なる場合は、従来どおり T に対する直列の計算を使います。

計算量: 係数数を \(n\) として \(O(n)\) 時間、追加空間は \(O(1)\) です。

依存ライブラリ

検証コード

コード

#ifndef CP_COMPONENTS_FPS_HORNER_HPP
#define CP_COMPONENTS_FPS_HORNER_HPP

#include "mint_arithmetic.hpp"
#include <array>
#include <cstring>

#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC push_options
#pragma GCC optimize("O3")
#endif

namespace fps_detail {
using mint_detail::Ops;
namespace horner_detail {
template<int K, class T>
inline T scalar(const T* a, std::size_t n, T x) {
	static_assert(K > 0);
	T step = 1;
	for (int i = 0; i < K; ++i) step *= x;
	std::array<T, K> sum{};
	std::size_t groups = n / K, tail = n % K;
	for (std::size_t i = 0; i < tail; ++i) sum[i] = a[groups * K + i];
	while (groups) {
		--groups;
		for (int i = 0; i < K; ++i) sum[i] = sum[i] * step + a[groups * K + i];
	}
	T result = 0;
	for (int i = K; i-- > 0;) result = result * x + sum[i];
	return result;
}

#if defined(__x86_64__) || defined(__i386__)
template<int K>
__attribute__((target("avx2"))) inline mint vector(const mint* a, std::size_t n, mint x) {
	static_assert(K == 1 || K == 2 || K == 4);
	using V = Ops::I8;
	static_assert(std::is_trivially_copyable_v<mint> && sizeof(mint) == sizeof(unsigned));
	mint step = x;
	for (int i = 0; i < 3; ++i) step *= step;
	const unsigned x8 = step.val();
	for (int i = 1; i < K; i *= 2) step *= step;
	const unsigned value = step.val();
	const V factor = _mm256_set1_epi32(value), quotient = _mm256_set1_epi32(Ops::quotient(value));
	V sum[K];
	for (auto& v : sum) v = _mm256_setzero_si256();
	constexpr std::size_t width = 8 * K;
	std::size_t groups = n / width, tail = n % width;
	if (tail) {
		alignas(32) mint block[width]{};
		std::memcpy(static_cast<void*>(block), a + groups * width, tail * sizeof(mint));
		for (int i = 0; i < K; ++i) std::memcpy(&sum[i], block + 8 * i, sizeof(V));
	} else if (groups) {
		--groups;
		for (int i = 0; i < K; ++i) std::memcpy(&sum[i], a + groups * width + 8 * i, sizeof(V));
	}
	while (groups) {
		--groups;
		for (int i = 0; i < K; ++i) {
			V coefficient;
			std::memcpy(&coefficient, a + groups * width + 8 * i, sizeof(V));
			sum[i] = Ops::add8(value == 1 ? sum[i] : Ops::fixed_mul8(sum[i], factor, quotient), coefficient);
		}
	}
	const V step8 = _mm256_set1_epi32(x8), quotient8 = _mm256_set1_epi32(Ops::quotient(x8));
	V joined = sum[K - 1];
	for (int i = K - 1; i-- > 0;) joined = Ops::add8(Ops::fixed_mul8(joined, step8, quotient8), sum[i]);
	alignas(32) mint word[8];
	std::memcpy(static_cast<void*>(word), &joined, sizeof(joined));
	mint result = 0;
	for (int i = 8; i-- > 0;) result = result * x + word[i];
	return result;
}
#endif
}

inline mint horner(const mint* a, std::size_t n, mint x) {
	if (!n) return mint(0);
	if (x == mint(0)) return a[0];
#if defined(__x86_64__) || defined(__i386__)
	if (n >= 32 && Ops::vectorized()) {
		if (n < 64) return horner_detail::vector<1>(a, n, x);
		return horner_detail::vector<4>(a, n, x);
	}
#endif
	if (n >= 32) return horner_detail::scalar<8>(a, n, x);
	mint result = 0;
	while (n) result = result * x + a[--n];
	return result;
}
}

#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC pop_options
#endif

#endif
Back to top page