cp-library

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

View the Project on GitHub qandori/cp-library

:heavy_check_mark: 多点評価の内部直接計算

components/fps_direct_eval.hpp · GitHub で見る

共通ページ

fps_detail::evaluate_direct(dst,f,n,points,m) は 多点評価 が使う内部処理です。通常は fps/multipoint_eval.hpp を include して multipoint_eval を利用してください。

関数 内容・返り値 計算量
fps_detail::evaluate_direct(dst,f,n,points,m) f の先頭 n 個を係数とする多項式を points の先頭 m 個で評価し、dst に書き込む \(O(nm)\)

f の先頭 n 個を係数として、points の先頭 m 個での評価値を dst に書き込みます。各配列は法 998244353 の内部 Mint ワード表現を使い、出力領域は入力領域と重ならないものとします。

この内部関数は x86 のみで定義され、AVX2 対応を実行時に確認してから呼び出す必要があります。公開関数は対応 CPU と対象サイズの場合だけこの経路を選び、非対応 CPU と Mac ARM では既存のスカラー処理を使います。

時間 \(O(nm)\)、出力領域を除く追加空間 \(O(1)\) です。

依存ライブラリ

検証コード

コード

#ifndef CP_COMPONENTS_FPS_DIRECT_EVAL_HPP
#define CP_COMPONENTS_FPS_DIRECT_EVAL_HPP

#include "mint_arithmetic.hpp"
#include <algorithm>
#include <cstddef>

namespace fps_detail {
using mint_detail::Ops;

#if defined(__x86_64__) || defined(__i386__)
__attribute__((target("avx2"))) inline void evaluate_direct(Ops::U *dst, const Ops::U *f,
		std::size_t n, const Ops::U *points, std::size_t m) {
	using I = Ops::I8;
	for (std::size_t i = 0; i < m; i += 32) {
		alignas(32) Ops::U tail[32] = {}, out[32];
		const std::size_t count = std::min<std::size_t>(32, m - i);
		if (count < 8) {
			for (std::size_t k = 0; k < count; ++k) {
				mint value = 0, x = mint::word(points[i + k]);
				for (std::size_t j = n; j-- > 0;) value = value * x + mint::word(f[j]);
				dst[i + k] = value.v;
			}
			break;
		}
		const Ops::U *xs = points + i;
		if (count < 32) {
			std::copy_n(xs, count, tail);
			xs = tail;
		}
		I x0 = Ops::load8(xs), x1 = Ops::load8(xs + 8);
		I x2 = Ops::load8(xs + 16), x3 = Ops::load8(xs + 24);
		I a0 = _mm256_setzero_si256(), a1 = a0, a2 = a0, a3 = a0;
		for (std::size_t j = n; j-- > 0;) {
			I c = _mm256_set1_epi32(f[j]);
			a0 = Ops::add8(Ops::mul8(a0, x0), c);
			a1 = Ops::add8(Ops::mul8(a1, x1), c);
			a2 = Ops::add8(Ops::mul8(a2, x2), c);
			a3 = Ops::add8(Ops::mul8(a3, x3), c);
		}
		Ops::U *output = count == 32 ? dst + i : out;
		Ops::store8(output, a0);
		Ops::store8(output + 8, a1);
		Ops::store8(output + 16, a2);
		Ops::store8(output + 24, a3);
		if (count < 32) std::copy_n(out, count, dst + i);
	}
}
#endif

}

#endif
Back to top page