FPS の内部演算

components/fps_operations.hpp · GitHub で見る

共通ページ

fps_detail::Operations<T> は FPS と独立した多項式関数で共有する内部処理です。通常は fps/fps.hpp や各関数のヘッダーから利用します。F は FormalPowerSeries<T>、W は内部係数型の FPS、base は vector<T>、Word は内部の 32 bit 整数です。

操作	内容・返り値	計算量
import_work(a) / export_work(a)	係数型を内部型へ／内部型から変換	\(O(n)\)
transform(n)	根を少なくとも n 個保持する共有 NTT プラン	拡張: \(O(n)\)、準備済み: \(O(1)\)
words(a)	係数配列の内部ワードへのポインタ	\(O(1)\)
transform_inplace<inverse>(a)	内部順序の NTT。逆変換は未正規化	\(O(n\log n)\)
ntt(a) / ntt_inverse(a) / intt(a)	順変換／未正規化の逆変換／正規化した逆変換	\(O(n\log n)\)
ntt_half(a) / ntt_half(a,n)	後半が 0 の係数列を変換	\(O(n\log n)\)
scale_words(a,c)	全係数を定数倍	\(O(n)\)
geometric_scale(a,first,step)	第 i 係数を first * step^i 倍	\(O(n)\)
batch_inverse(dst,src)	各非零要素の逆元を一括計算	\(O(n)\)
pointwise(r,a,b,n)	先頭 n 個の要素ごとの積	\(O(n)\)
accumulate_products<count,first>(dst,a,b,n)	count 組の先頭 n 個の要素ごとの積を加算。first が真なら上書き、偽なら dst へ加算	\(O(n)\)
accumulate_products8<count,first>(dst,a,b,n)	accumulate_products と同じ演算を AVX2 を有効にして実行	\(O(n)\)
join_windows(blocks,length,count)	連続する count 個の変換済みブロックを接続。前半に直前のブロックを加え、後半から引く	\(O(\mathrm{length}\cdot\mathrm{count})\)
join_windows_strided(blocks,length,count,pitch)	要素数 pitch の間隔で並べたブロックを接続	\(O(\mathrm{length}\cdot\mathrm{count})\)
inverse_extend(h,a,cache)	キャッシュした変換を使って逆数の精度を倍化	\(O(n\log n)\)
dot(a,b,n)	逆向きの添字を使う内積	\(O(n)\)
mod_sqrt(a)	法 998244353 の平方根。非平方剰余は 0	\(O(1)\)
inverse_ntt_roots(n) / ntt_roots(n)	内部順序の逆根／根のキャッシュ	拡張: \(O(n)\)、準備済み: \(O(1)\)
inverses(n)	整数 0～n の逆元のキャッシュ。0 の項は 0	拡張: \(O(n)\)、準備済み: \(O(1)\)
trim(a)	末尾の零係数を除去	\(O(n)\)

accumulate_products は \(1\le\mathrm{count}\le7\) とし、a・b は各 count 個の配列ポインタを渡します。各配列と dst に n 要素を用意し、内部 Montgomery ワードとして a[k][i] は \([0,p)\)、b[k][i] は \([0,2p)\) に収めます。ここで \(p=998244353\) です。first = false のときは元の dst も \([0,2p)\) にします。結果も \([0,2p)\) です。n が 8 以上なら 8 の倍数にします。accumulate_products8 は x86 のみで定義され、n を 8 の倍数とし、AVX2 対応を確認してから呼び出します。

join_windows は length を 2 以上の 2 冪、count を 1 以上とします。各ブロックは、前半に係数を置き、後半を 0 とした長さ length の NTT です。全ブロックで count * length 要素を用意します。join_windows_strided では、ブロック先頭間の要素数 pitch を length 以上とし、count * pitch 要素を用意します。接続前の直前のブロックを参照するため、処理は後ろから行います。

係数数を \(n\) として、変換・inverse_extend は \(O(n\log n)\)、その他の配列操作は \(O(n)\) です。根・逆元のキャッシュ拡張は追加要素数に比例し、取得のみなら \(O(1)\)、words は \(O(1)\) です。mod_sqrt は固定法に対して \(O(1)\)、追加空間・キャッシュは \(O(n)\) です。

accumulate_products・accumulate_products8 は時間 \(O(n)\)、join_windows は時間 \(O(\mathrm{length}\cdot\mathrm{count})\) です。これらの追加空間は \(O(1)\) です。

依存ライブラリ

• modint998244353 の共通剰余演算
• NTT（998244353）（高速版）
• modint998244353

検証コード

• test/convolution.test.cpp
• test/fps_bostan_mori.test.cpp
• test/fps_compose.test.cpp
• test/fps_compose_large.test.cpp
• test/fps_compositional_inverse.test.cpp
• test/fps_compositional_inverse_large.test.cpp
• test/fps_divmod.test.cpp
• test/fps_exp.test.cpp
• test/fps_exp_sparse.test.cpp
• test/fps_interpolate.test.cpp
• test/fps_inv.test.cpp
• test/fps_inv_sparse.test.cpp
• test/fps_linear_recurrence.test.cpp
• test/fps_log.test.cpp
• test/fps_log_sparse.test.cpp
• test/fps_pow.test.cpp
• test/fps_pow_sparse.test.cpp
• test/fps_rational_coefficients.test.cpp
• test/fps_sqrt.test.cpp
• test/fps_sqrt_sparse.test.cpp
• test/fps_taylor_shift.test.cpp
• test/multipoint_eval.test.cpp
• test/polynomial_inv_mod.test.cpp
• test/power_projection.test.cpp

コード

#ifndef CP_COMPONENTS_FPS_OPERATIONS_HPP
#define CP_COMPONENTS_FPS_OPERATIONS_HPP

#include "ntt.hpp"
#include <algorithm>
#include <cassert>
#include <numeric>
#include <type_traits>
#include <vector>

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

template<class T> struct FormalPowerSeries;

namespace fps_detail {

template<class T, class Work = mint>
struct Operations {
	using F = FormalPowerSeries<T>;
	using W = FormalPowerSeries<Work>;
	using base = std::vector<T>;
	using Word = NTT::U;

	static W import_work(const base &a) {
		W r(a.size());
		for (int i = 0; i < static_cast<int>(a.size()); ++i) r[i] = mint::raw(a[i].val());
		return r;
	}
	static F export_work(const W &a) {
		F r(a.size());
		for (int i = 0; i < static_cast<int>(a.size()); ++i) r[i] = T::raw(a[i].val());
		return r;
	}
	static NTT &transform(int n) { return ntt_plan(n); }
	static Word *words(base &a) {
		static_assert(sizeof(T) == sizeof(Word) && std::is_trivially_copyable_v<T>);
		return reinterpret_cast<Word *>(a.data());
	}
	static const Word *words(const base &a) { return reinterpret_cast<const Word *>(a.data()); }
	template <bool inverse>
	static void transform_inplace(base &a) {
		int n = a.size();
		assert(n && !(n & (n - 1)) && n <= (1 << 23));
		auto &t = transform(std::max(1, n / 2));
		t.template fft<inverse>(words(a), n);
		if constexpr (!std::is_same_v<T, mint>) t.normalize(words(a), n);
	}
	static void ntt(base &a) { transform_inplace<false>(a); }
	static void ntt_half(Word *a, int n) {
		std::copy_n(a, n / 2, a + n / 2);
		transform(std::max(1, n / 2)).template fft<false>(a, n, n / 2);
	}
	static void ntt_half(base &a) { ntt_half(words(a), a.size()); }
	static void ntt_inverse(base &a) { transform_inplace<true>(a); }
	static void scale_words(base &a, T factor) {
		if (a.empty()) return;
		if constexpr (std::is_same_v<T, mint>) {
			auto &t = transform(1);
			t.scale(words(a), a.size(), factor.val());
		} else
			for (T &x : a) x *= factor;
	}
	static void geometric_scale(base &a, T first, T step) {
		if constexpr (std::is_same_v<T, mint>) NTT::geometric(words(a), a.size(), first.v, step.v);
		else
			for (T &x : a) {
				x *= first;
				first *= step;
			}
	}
	static void batch_inverse(base &dst, const base &src) {
		int n = src.size();
#if defined(__x86_64__) || defined(__i386__)
		if constexpr (std::is_same_v<T, mint>) {
			if (n >= 128 && NTT::vectorized()) {
				NTT::batch_inverse8(words(dst), words(src), n);
				return;
			}
		}
#endif
		base prefix(n);
		T p = 1;
		for (int i = 0; i < n; ++i) prefix[i] = p, p *= src[i];
		assert(p != T(0));
		p = p.inv();
		for (int i = n - 1; i >= 0; --i) {
			dst[i] = prefix[i] * p;
			p *= src[i];
		}
	}
	static void pointwise(base &r, const base &a, const base &b, int n) {
		if constexpr (std::is_same_v<T, mint>) NTT::product(words(r), words(a), words(b), n);
		else
			for (int i = 0; i < n; ++i) r[i] = a[i] * b[i];
	}
	template <int count, bool first>
	static void accumulate_products(Word *dst, const Word *const *a, const Word *const *b, int n) {
		static_assert(count >= 1 && count <= 7);
		for (int i = 0; i < n; ++i) {
			Word value = 0;
			for (int k = 0; k < count; ++k) value = NTT::add(value, NTT::mul(a[k][i], b[k][i]));
			if constexpr (!first) value = NTT::add(value, dst[i]);
			dst[i] = value;
		}
	}
#if defined(__x86_64__) || defined(__i386__)
	template<int count, bool first>
	__attribute__((target("avx2"))) static void accumulate_products8(Word *dst, const Word *const *a, const Word *const *b, int n) {
		accumulate_products<count, first>(dst, a, b, n);
	}
#endif
	static void join_windows(Word *blocks, int length, int count) {
		join_windows_strided(blocks, length, count, length);
	}
	static void join_windows_strided(Word *blocks, int length, int count, int pitch) {
		for (int i = count - 1; i > 0; --i) {
			Word *a = blocks + i * pitch;
			const Word *b = a - pitch;
			for (int j = 0; j < length / 2; ++j) a[j] = NTT::add(a[j], b[j]);
			for (int j = length / 2; j < length; ++j) a[j] = NTT::sub(a[j], b[j]);
		}
	}
	static void inverse_extend(F &h, const base &a, const base &cache) {
		const int m = static_cast<int>(cache.size());
		const T im = T::raw(T::mod() - (T::mod() - 1) / m);
		base b(m);
		pointwise(b, a, cache, m);
		ntt_inverse(b);
		std::fill(b.begin(), b.begin() + m / 2, T(0));
		ntt(b);
		pointwise(b, b, cache, m);
		ntt_inverse(b);
		h.resize(m);
		for (int i = m / 2; i < m; ++i) h[i] = -b[i] * im * im;
	}
	static T dot(const T *a, const T *b, int n) {

#if defined(__x86_64__) || defined(__i386__)
		if constexpr (std::is_same_v<T, mint>) {
			if (n >= 64 && NTT::vectorized())
				return T::word(NTT::reverse_dot8(reinterpret_cast<const Word *>(a), reinterpret_cast<const Word *>(b), n));
		}
#endif
		unsigned int result = 0;
		for (int i = 0; i < n; i += 16) {
			unsigned long long sum = 0;
			for (int j = i; j < std::min(n, i + 16); ++j)
				sum += (unsigned long long)a[j].val() * b[-j].val();
			result += sum % T::mod();
			if (result >= T::mod()) result -= T::mod();
		}
		return T::raw(result);
	}

	static void intt(base &a) {
		ntt_inverse(a);
		const T inv = T::raw(T::mod() - (T::mod() - 1) / static_cast<int>(a.size()));
		scale_words(a, inv);
	}
	static T mod_sqrt(T a) {
		if (a == T(0)) return 0;
		T t = a.pow(59), r = t * a, c = T(3).pow(119);
		t *= r;
		int m = 23;
		while (t != T(1)) {
			int i = 0;
			for (T u = t; u != T(1) && i < m; u *= u) ++i;
			if (i == m) return 0;
			T b = c.pow(1 << (m - i - 1));
			r *= b;
			c = b * b;
			t *= c;
			m = i;
		}
		return r;
	}
	template<bool Inverse>
	static const base &roots(int n) {
		static base values{1};
		while (static_cast<int>(values.size()) < n) {
			const int m = static_cast<int>(values.size());
			T step = T(3).pow((T::mod() - 1) / (2 * m));
			if constexpr (Inverse) step = step.inv();
			values.resize(2 * m);
			for (int i = 0; i < m; ++i) values[m + i] = values[i] * step;
		}
		return values;
	}
	static const base &inverse_ntt_roots(int n) { return roots<true>(n); }
	static const base &ntt_roots(int n) { return roots<false>(n); }
	static const std::vector<T> &inverses(int n) {
		assert(unsigned(n) < unsigned(T::mod()));
		static std::vector<T> inv{0, 1};
#if defined(__x86_64__) || defined(__i386__)
		if constexpr (std::is_same_v<T, mint>) {
			int first = inv.size(), count = (n + 1 - first) / 8 * 8;
			if (count >= 256 && NTT::vectorized()) {
				inv.reserve(n + 1);
				inv.resize(first + count);
				NTT::inverse_sequence8(reinterpret_cast<Word *>(inv.data()) + first, first, count);
			}
		}
#endif
		while (static_cast<int>(inv.size()) <= n) {
			const int i = static_cast<int>(inv.size());
			inv.push_back(-inv[T::mod() % i] * (T::mod() / i));
		}
		return inv;
	}
	static void trim(F &a) {
		while (!a.empty() && a.back() == T(0)) a.pop_back();
	}
};

}

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

#endif

Back to top page