cp-library

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

View the Project on GitHub qandori/cp-library

:heavy_check_mark: 素数 mod 上の乗法畳み込み(高速版)

convolution/multiplicative_convolution.hpp · GitHub で見る

全体サイト(256 KiB向け) · 高速実装

#include "lib/convolution/multiplicative_convolution.hpp"

std::vector<mint> a{1, 2, 3}, b{4, 5, 6};
auto c = multiplicative_convolution(a, b); // {35,28,27}
操作 内容・返り値 計算量
multiplicative_convolution(a,b) 同じ素数長 P の列から \(\displaystyle c_k=\sum_{ij\bmod P=k}a_ib_j\) を返す \(O(M(P)+P)\)+原始根探索

添字 i,j の積を P で割った余りを出力先にします。係数の計算は入力型の加算・乗算です。例は長さ P=3 なので、a[2]*b[2] は 2*2 % 3 = 1 より c[1] に加えます。c は入力と同じ型の長さ P の配列です。

入力は同じ長さの素数 \(P\) 個とし、\(\displaystyle c_k = \sum_{ij\bmod P=k} a_i b_j\) を返します。係数型の条件は 通常の畳み込み と同じです。

原始根探索後、通常の長さ \(P\) の畳み込み時間を \(M(P)\) として、時間 \(O(M(P)+P)\)、追加空間 \(O(P)\) です。原始根探索では \(P-1\) の素因数分解と候補ごとの \(O(\log^{2} P)\) の判定を行います。

依存ライブラリ

検証コード

コード

#ifndef MULTIPLICATIVE_CONVOLUTION_HPP
#define MULTIPLICATIVE_CONVOLUTION_HPP

#include "../components/modular_power_sequence.hpp"
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <limits>
#include <vector>

#include "../math/factorize.hpp"
#include "../components/convolution_dispatch.hpp"

template <class T>
std::vector<T> multiplicative_convolution(const std::vector<T> &a, const std::vector<T> &b) {
	if (a.empty() || b.empty()) return {};
	assert(a.size() == b.size() && a.size() <= std::numeric_limits<std::uint32_t>::max() && is_prime(a.size()));
	std::uint32_t p = a.size();
	auto factors = factorize(p - 1);
	factors.erase(std::unique(factors.begin(), factors.end()), factors.end());
	auto power = [=](std::uint64_t x, std::uint32_t k) {
		std::uint64_t result = 1;
		for (; k; k >>= 1, x = x * x % p)
			if (k & 1) result = result * x % p;
		return result;
	};
	std::uint32_t root = 1;
	while (std::any_of(factors.begin(), factors.end(), [&](auto q) { return power(root, (p - 1) / q) == 1; })) ++root;
	std::vector<T> x(p - 1), y(&a == &b ? 0 : p - 1), result(p, T(0));
	auto powers = modular_power_sequence_detail::generate(p - 1, root, p);
	for (std::uint32_t i = 0; i < p - 1; ++i) {
		x[i] = a[powers[i]];
		if (&a != &b) y[i] = b[powers[i]];
	}

	auto product = &a == &b ? convolution(x, x) : convolution(x, y);
	for (std::size_t i = p - 1; i < product.size(); ++i) product[i - (p - 1)] += product[i];
	for (std::uint32_t i = 0; i < p - 1; ++i) result[powers[i]] = product[i];
	T sum_a(0), sum_b(0);
	for (std::size_t i = 1; i < a.size(); ++i) sum_a += a[i];
	for (const T &v : b) sum_b += v;
	result[0] = a[0] * sum_b;
	result[0] += sum_a * b[0];
	return result;
}

#endif
Back to top page