FPS の非対称畳み込み

components/fps_convolution_unbalanced.hpp · GitHub で見る

共通ページ

Linux x86_64 / GCC の AVX2 用内部部品です。公開の fps::multiply と convolution から利用します。

操作	内容	計算量
fps_detail::multiply_detail::unbalanced::fits(n,m)	短辺が 32 より大きく、長辺が 32768 以上かつ短辺の 32 倍以上か判定	\(O(1)\)
fps_detail::multiply_detail::unbalanced::multiply(a,b)	読み取り専用の二つの係数列の積を返す	\(O((n+m)\log m)\)

ここで \(n\) は長辺、\(m\) は短辺です。multiply の呼び出し側は fits(a.size(),b.size())、AVX2 対応、両 span が生存中の Mint 要素を指すこと、非空入力、入力の内部 Montgomery 値が \([0,2p)\)、出力長が \(2^{23}\) 以下であることを保証します。公開 dispatch は fits を満たす場合だけこの経路を使います。

短辺の通常剰余の変換を一度作り、長辺を幅 \(L-m+1\) のブロックに分けます。\(L=8\operatorname{bitceil}(m)\) とし、各ブロックの循環積を線形積として復元し、隣接ブロックの重なり \(m-1\) 項を加算します。長辺は Montgomery 表現のまま計算し、返す係数も正規化した Montgomery 表現です。

入力を変更せず、公開入力に alignment や padding を要求しません。作業配列は Arena が所有し、出力は独立した vector<Mint> の生存中の要素として構築します。同じ入力を二回渡しても安全です。追加の作業空間は \(O(m)\)、返す配列は \(O(n+m)\) です。

依存ライブラリ

• FPS のブロック変換
• FPS の結果配列への一括変換
• FPS の作業領域
• modint998244353 の共通剰余演算
• modint998244353

検証コード

• test/convolution.test.cpp
• test/cyclic_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/multiplicative_convolution.test.cpp
• test/multipoint_eval.test.cpp
• test/multivariate_convolution.test.cpp
• test/ordered_convolution.test.cpp
• test/polynomial_inv_mod.test.cpp
• test/power_projection.test.cpp

コード

#ifndef CP_COMPONENTS_FPS_CONVOLUTION_UNBALANCED_HPP
#define CP_COMPONENTS_FPS_CONVOLUTION_UNBALANCED_HPP

#include "fps_blocked.hpp"
#include <algorithm>
#include <bit>
#include <cstddef>
#include <span>
#include <vector>

#if defined(__linux__) && defined(__x86_64__) && defined(__GNUC__) && !defined(__clang__)
#pragma GCC push_options
#pragma GCC optimize("O3")
_Pragma("GCC target(\"avx2\")")
namespace fps_detail::multiply_detail::unbalanced {
using namespace blocked_poly;
inline bool fits(size_t n, size_t m) {
	const auto shorter = std::min(n, m), longer = std::max(n, m);
	return shorter > 32 && longer >= 32768 && shorter <= longer / 32;
}
inline std::vector<mint> multiply(std::span<const mint> a, std::span<const mint> b) {
	const size_t len = 8 * std::bit_ceil(std::min(a.size(), b.size()));
	const auto& x = a.size() >= b.size() ? a : b;
	const auto& y = a.size() >= b.size() ? b : a;
	const size_t n = x.size(), m = y.size(), stride = len - m + 1;
	const int lg = std::countr_zero(len);
	Arena arena(Transform::words(lg) + 3 * Arena::footprint(len));
	const Transform t(arena, lg);
	auto bt = arena.take(len), in = arena.take(len), out = arena.take(len);
	for (size_t i = 0; i < m; ++i) in[i] = y[i].val();
	t.forward(in.first(m), 0, bt);
	std::vector<mint> result;
	fps_detail::prepare_fps_output(result, n + m - 1);
	for (size_t p = 0; p < n; p += stride) {
		const size_t count = std::min(stride, n - p);
		for (size_t i = 0; i < count; ++i) {
			const unsigned v = x[p + i].v;
			in[i] = std::min(v, v - mint::mod());
		}
		t.cyclic_product(in.first(count), 0, out, bt);
		const size_t overlap = p == 0 ? 0 : m - 1;
		for (size_t i = 0; i < overlap; ++i) {
			const unsigned v = result[p + i].v + out[i];
			result[p + i] = mint::word(std::min(v, v - mint::mod()));
		}
		fps_detail::append_fps_words<true>(result, out.data() + overlap, count + m - 1 - overlap);
	}
	return result;
}
}
#pragma GCC pop_options
#endif
#endif

Back to top page