cp-library

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

View the Project on GitHub qandori/cp-library

:heavy_check_mark: 多変数の切り捨て畳み込み(軽量実装)

convolution/multivariate_convolution_light.hpp · GitHub で見る

共通ページ · 軽量実装

light_convolution 名前空間で呼び出します。通常版と同時に include できます。

#include "lib/convolution/multivariate_convolution_light.hpp"

std::vector<mint> a{1, 2, 3, 4, 5, 6}, b{7, 8, 9, 10, 11, 12};
auto c = light_convolution::multivariate_convolution(a, b, {2, 3}); // {7,22,30,80,73,182}
操作 内容・返り値 計算量
light_convolution::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)\) です。

依存ライブラリ

検証コード

コード

#ifndef MULTIVARIATE_CONVOLUTION_LIGHT_HPP
#define MULTIVARIATE_CONVOLUTION_LIGHT_HPP

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

#include "convolution_light.hpp"

namespace light_convolution {
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 = light_convolution::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());
			auto transform = [](std::vector<mint> &v, bool inverse) { fps_light_detail::transform(v, inverse); };
			auto product = multivariate_convolution_detail::Core::multiply(x, y, dims, transform, fps_light_detail::multiply);
			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 = light_convolution::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