多変数の切り捨て畳み込み（高速版）

convolution/multivariate_convolution.hpp · GitHub で見る

全体サイト（256 KiB向け） · 高速実装

軽量実装（64 KiB）はこちら

#include "lib/convolution/multivariate_convolution.hpp"

std::vector<mint> a{1, 2, 3, 4, 5, 6}, b{7, 8, 9, 10, 11, 12};
auto c = multivariate_convolution(a, b, {2, 3}); // {7,22,30,80,73,182}

操作	内容・返り値	計算量
multivariate_convolution(a,b,dimensions)	各変数の次数上限を dimensions で指定して、切り捨てた積を返す	998244353: \(O(K+DN\log N+D^2N)\)、その他: \(O(K+M(ND))\)

例の {2,3} は、\(x\) の指数を 0〜1、\(y\) の指数を 0〜2 に制限する指定です。係数を \(1,x,y,xy,y^2,xy^2\) の順に並べた長さ 6 の配列を渡し、同じ順序・長さで返します。積の \(x^2\) 以上や \(y^3\) 以上の項は捨てます。

多項式の積を \((x_0^{d_0},\ldots,x_s^{d_s})\) で切り捨てます。各 \(d > 0\)、入力長は \(\displaystyle N = \prod_{j=0}^{s} d_j\) とし、添字は \(i_0 + d_0 i_1 + d_0 d_1 i_2 + \cdots\) の順です。

入力の次元数を \(K\)、1 より大きい次元数を \(D\)、通常の長さ \(L\) の畳み込み時間を \(M(L)\) とします。mod 998244353 では \(N \le 2^{22}\)、時間 \(O(K + DN \log N + D^{2}N)\)、その他は時間 \(O(K + M(ND))\) です。係数型・変換長の条件は 通常の畳み込み と同じで、追加空間はいずれも \(O(ND)\) です。\(N=1\) は時間 \(O(K)\)、追加空間 \(O(1)\) です。

依存ライブラリ

• 畳み込みの型別振り分け
• FFT
• FPS のブロック変換
• FPS の融合畳み込み
• FPS の定数係数畳み込み
• FPS の非対称畳み込み
• FPS の結果配列への一括変換
• FPS の作業領域
• modint998244353 の共通剰余演算
• modint998244353 の共有畳み込み
• 多変数畳み込みの共通実装
• 小さい次元の多変数畳み込み
• NTT（998244353）（高速版）
• NTT（法指定）
• 実数・複素数の畳み込み
• 任意 mod・整数の畳み込み
• 一般の環の畳み込み
• modint998244353

検証コード

• test/multivariate_convolution.test.cpp

コード

#ifndef MULTIVARIATE_CONVOLUTION_HPP
#define MULTIVARIATE_CONVOLUTION_HPP

#include <cassert>
#include <cstddef>
#include <vector>
#include "../components/multivariate_convolution_core.hpp"
#include "../components/multivariate_ranked.hpp"

#include "../components/convolution_dispatch.hpp"
#include "../components/ntt.hpp"

template <class T>
std::vector<T> multivariate_convolution(const std::vector<T> &a, const std::vector<T> &b, const std::vector<int> &dimensions) {
	std::size_t size = 1;
	std::vector<std::size_t> bases;
	for (int d : dimensions) {
		assert(d > 0 && size <= std::size_t(-1) / d);
		if (d == 1) continue;
		size *= d;
		bases.push_back(size);
	}
	assert(a.size() == size && b.size() == size);
	if (bases.empty()) return {a[0] * b[0]};
	if (bases.size() == 1) {
		auto result = convolution(a, b);
		result.resize(size);
		return result;
	}
	int colors = bases.size();

	if constexpr (requires { T::mod(); T::raw(0); a[0].val(); }) {
		if constexpr (T::mod() == 998244353) {
			assert(size <= std::size_t(1) << 22);
			std::vector<mint> x(size), y(size);
			std::vector<int> dims;
			for (int d : dimensions) if (d > 1) dims.push_back(d);
			for (std::size_t i = 0; i < size; ++i) x[i] = mint::raw(a[i].val()), y[i] = mint::raw(b[i].val());
			std::unique_ptr<fps_detail::NTT> ntt;
			auto transform = [&](std::vector<mint> &v, bool inverse) {
				using NTT = fps_detail::NTT;
				if (!ntt) ntt = std::make_unique<NTT>(v.size() / 2, false);
				auto *p = reinterpret_cast<NTT::U *>(v.data());
				if (inverse) ntt->fft<true>(p, v.size()), ntt->scale(p, v.size(), 998244353 - 998244352 / v.size());
				else ntt->fft<false>(p, v.size());
			};
			auto product = [&] {
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
				int degree = 0;
				bool small = true;
				for (int d : dims) degree += d - 1, small &= d <= 3;
				if (small && degree < 24 && std::count(dims.begin(), dims.end(), 2) >= 3 && multivariate_convolution_detail::Core::available())
					return multivariate_ranked_detail::multiply(x, y, dims);
#endif
				return multivariate_convolution_detail::Core::multiply(x, y, dims, transform, convolution<mint>);
			}();
			std::vector<T> result;
			result.reserve(size);
			for (auto v : product) result.push_back(T::raw(v.val()));
			return result;
		}
	}
	std::vector<int> rank(size);
	for (std::size_t i = 0; i < size; ++i) {
		std::size_t value = 0;
		for (std::size_t base : bases) value += i / base;
		rank[i] = value % colors;
	}
	std::size_t stride = 2 * colors - 1;
	assert(size <= std::size_t(-1) / stride / 2);
	std::vector<T> x(size * stride, T(0)), y(size * stride, T(0));
	for (std::size_t i = 0; i < size; ++i) x[i * stride + rank[i]] = a[i], y[i * stride + rank[i]] = b[i];
	auto product = convolution(x, y);
	std::vector<T> result(size, T(0));
	for (std::size_t i = 0; i < size; ++i) {
		result[i] = product[i * stride + rank[i]];
		if (rank[i] + colors < int(stride)) result[i] += product[i * stride + rank[i] + colors];
	}
	return result;
}

#endif

Back to top page