cp-library

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

View the Project on GitHub qandori/cp-library

:heavy_check_mark: 一般の環の畳み込み

convolution/convolution_ring.hpp · GitHub で見る

共通ページ

#include "lib/convolution/convolution_ring.hpp"

std::vector<long long> a{1, 2, 3}, b{4, 5};
auto c = convolution_ring(a, b); // {4,13,22,15}
操作 内容・返り値 計算量
convolution_ring(a,b) 一般の環で通常の畳み込みを計算 \(O(L^{\log_2 3})\)

\(\displaystyle c_k = \sum_{i+j=k} a_i b_j\) を Karatsuba 法で計算します。係数型には T(0)・+=・-=・* が必要です。整数型では中間値も型の範囲内とします。

vector<T> を 2 個渡すと vector<T> で積を返します。多項式は定数項から並べます。両方が非空なら長さは a.size()+b.size()-1、どちらかが空なら空配列です。入力は変更しません。NTT に必要な逆元や原始根を持たない型にも使えます。

\(L = \lvert a\rvert + \lvert b\rvert\) として、時間 \(O(L^{\log_2 3})\)、追加空間 \(O(L)\) です。

検証コード

コード

#ifndef CONVOLUTION_RING_HPP
#define CONVOLUTION_RING_HPP

#include <algorithm>
#include <cstddef>
#include <span>
#include <vector>

namespace convolution_ring_detail {
template <class T>
std::vector<T> multiply(std::span<const T> a, std::span<const T> b) {
	if (a.empty() || b.empty()) return {};
	std::vector<T> result(a.size() + b.size() - 1, T(0));
	if (std::min(a.size(), b.size()) <= 32) {
		for (std::size_t i = 0; i < a.size(); ++i)
			for (std::size_t j = 0; j < b.size(); ++j) result[i + j] += a[i] * b[j];
		return result;
	}
	std::size_t half = std::max(a.size(), b.size()) / 2;
	auto al = a.first(std::min(a.size(), half)), ar = a.subspan(al.size());
	auto bl = b.first(std::min(b.size(), half)), br = b.subspan(bl.size());
	auto low = multiply(al, bl), high = multiply(ar, br);
	std::vector<T> x(std::max(al.size(), ar.size()), T(0)), y(std::max(bl.size(), br.size()), T(0));
	for (std::size_t i = 0; i < al.size(); ++i) x[i] += al[i];
	for (std::size_t i = 0; i < ar.size(); ++i) x[i] += ar[i];
	for (std::size_t i = 0; i < bl.size(); ++i) y[i] += bl[i];
	for (std::size_t i = 0; i < br.size(); ++i) y[i] += br[i];
	auto middle = multiply<T>(x, y);
	for (std::size_t i = 0; i < low.size(); ++i) result[i] += low[i], middle[i] -= low[i];
	for (std::size_t i = 0; i < high.size(); ++i) result[i + 2 * half] += high[i], middle[i] -= high[i];
	for (std::size_t i = 0; i < middle.size() && i + half < result.size(); ++i) result[i + half] += middle[i];
	return result;
}
}

template <class T>
std::vector<T> convolution_ring(const std::vector<T> &a, const std::vector<T> &b) {
	return convolution_ring_detail::multiply<T>(a, b);
}

#endif
Back to top page