任意 mod 数論変換 (任意 mod NTT)
(Convolution/NumberTheoreticTransform_AnyMod.hpp)

View this file on GitHub
Last update: 2021-06-05 13:56:14+09:00
Include: #include "Convolution/NumberTheoreticTransform_AnyMod.hpp"
Link: https://tkmst201.github.io/Library/Convolution/NumberTheoreticTransform_AnyMod.hpp

概要

数論変換 (NTT) に Garner を組み合わせて、法 $M$ が任意に取れるようになったものです。ただし、積の各係数において $\bmod{M}$ を取る前の値がある値以下に収まっている必要があります。
$2$ つの整数係数多項式 $f(x), g(x)$ に対して、積 $f(x)g(x) \bmod{M}$ を計算します。計算量は $\Theta(N\log{N})$ (積の次数を $N$ ) です。

静的メンバ関数

std::vector<T> multiply(const std::vector<T> & a, const std::vector<T> & b)

$f(x) := \sum_{i=0}^{|a| - 1} a[i] x^i \bmod{MOD}$, $g(x) := \sum_{i=0}^{|b| - 1} b[i] x^i \bmod{MOD}$ として、積 $f(x)g(x) \equiv \sum_{i=0}^{|a| + |b| - 2} c[i] x^i \pmod{MOD}$ となるような大きさ $|a| + |b| - 1$ の列 $c[i] \bmod{MOD}$ 返します。
$a, b$ いずれかが空である場合は、空の列を返します。

制約

MOD $> 0$
$a_i, b_j \geq 0$
多項式 $a, b$ の係数の最大値を $Q$ として、$\max(|A|, |B|) * Q^2 < \prod_i m_i$ となるように $m$ が選ばれている
T は int, long long int

計算量

法 $m$ の個数を $K$ 、$N$ を $|a| + |b| - 1$ 以上の最小の $2$ 冪とします。

$\mathcal{O}(N\log{N} + (|a| + |b|) K(K + \log{\max{m_i}}))$

使用例

#include <bits/stdc++.h>
#include "Convolution/NumberTheoreticTransform_AnyMod.hpp"
using namespace std;

int main() {
	vector<int> a({0, 1, 2, 3}), b({2, 3, 4});
	auto c = NumberTheoreticTransform_AnyMod<14>::multiply(a, b);
	cout << "size = " << c.size() << endl; // 6
	for (int i = 0; i < c.size(); ++i) cout << c[i] << " "; cout << endl; // 0 2 7 2 3 12
	//   2 3 4
	//     4 6 8
	//       6 9 12
	// ==============
	// 0 2 7 16 17 12
	// 0 2 7  2  3 12
	
	vector<int> a2({4}), b2({3, 2});
	auto c2 = NumberTheoreticTransform_AnyMod<7>::multiply(a2, b2);
	cout << "size = " << c2.size() << endl; // 2
	for (int i = 0; i < c2.size(); ++i) cout << c2[i] << " "; cout << endl; // 5 1
	// 12 
	//    8
	// ==============
	// 12 8
	//  5 1
	
	vector<int> a3, b3({1, 2, 3});
	auto c3 = NumberTheoreticTransform_AnyMod<100>::multiply(a3, b3);
	cout << "size = " << c3.size() << endl; // 0
}

参考

2020/07/27: https://math314.hateblo.jp/entry/2015/05/07/014908

Depends on

Verified with

Test/NumberTheoreticTransform_AnyMod.test.cpp

Code

#ifndef INCLUDE_GUARD_NUMBER_THEORETIC_TRANSFORM_ANY_MOD_HPP
#define INCLUDE_GUARD_NUMBER_THEORETIC_TRANSFORM_ANY_MOD_HPP

#include "Convolution/NumberTheoreticTransform.hpp"
#include "Mathematics/garner.hpp"

#include <vector>
#include <cstdint>
#include <vector>
#include <type_traits>

/**
 * @brief https://tkmst201.github.io/Library/Convolution/NumberTheoreticTransform_AnyMod.hpp
 */
template<int MOD>
struct NumberTheoreticTransform_AnyMod {
	static_assert(MOD > 0);
	
private:
	using calc_type = long long;
	using uint32 = std::uint32_t;
	
public:
	template<typename T>
	static std::vector<T> multiply(const std::vector<T> & a, const std::vector<T> & b) {
		static_assert(std::is_integral<T>::value);
		static_assert(std::is_signed<T>::value);
		for (uint32 i = 0; i < a.size(); ++i) assert(a[i] >= 0);
		for (uint32 i = 0; i < b.size(); ++i) assert(b[i] >= 0);
		std::vector<T> m;
		auto ntt1_res = NumberTheoreticTransform<1'224'736'769, 3>::multiply(a, b);
		m.emplace_back(1'224'736'769);
		auto ntt2_res = NumberTheoreticTransform<469'762'049, 3>::multiply(a, b);
		m.emplace_back(469'762'049);
		auto ntt3_res = NumberTheoreticTransform<167'772'161, 3>::multiply(a, b);
		m.emplace_back(167'772'161);
		// auto ntt4_res = NumberTheoreticTransform<998'244'353, 3>::multiply(a, b);

		// m.emplace_back(998'244'353);

		
		std::vector<T> c(m.size());
		std::vector<T> res(ntt1_res.size());
		for (uint32 i = 0; i < res.size(); ++i) {
			c[0] = ntt1_res[i];
			c[1] = ntt2_res[i];
			c[2] = ntt3_res[i];
			// c[3] = ntt4_res[i];

			res[i] = tk::garner<T, calc_type>(c, m, MOD);
		}
		return res;
	}
};

#endif // INCLUDE_GUARD_NUMBER_THEORETIC_TRANSFORM_ANY_MOD_HPP

#line 1 "Convolution/NumberTheoreticTransform_AnyMod.hpp"



#line 1 "Convolution/NumberTheoreticTransform.hpp"



#line 1 "Mathematics/mod_pow_inv.hpp"



#include <cassert>
#include <type_traits>

/**
 * @brief https://tkmst201.github.io/Library/Mathematics/mod_pow_inv.hpp
 */
namespace tk {
template<typename T>
constexpr T mod_pow(T x, T n, T m) noexcept {
	static_assert(std::is_integral<T>::value);
	assert(m > 0);
	assert(n >= 0);
	x = x % m + (x >= 0 ? 0 : m);
	T res = 1 % m;
	while (n > 0) {
		if (n & 1) res = res * x % m;
		x = x * x % m;
		n >>= 1;
	}
	return res;
}

template<typename T>
constexpr T mod_inv(T x, T m) noexcept {
	static_assert(std::is_integral<T>::value);
	static_assert(std::is_signed<T>::value);
	assert(m > 0);
	x = x % m + (x >= 0 ? 0 : m);
	T x1 = 1, y = m, y1 = 0;
	while (y > 0) {
		const T q = x / y;
		T tmp = x - q * y; x = y; y = tmp;
		tmp = x1 - q * y1; x1 = y1; y1 = tmp;
	}
	assert(x == 1);
	if (x1 == m) x1 = 0;
	if (x1 < 0) x1 += m;
	return x1;
}
} // namespace tk



#line 5 "Convolution/NumberTheoreticTransform.hpp"

#include <vector>
#include <utility>
#line 9 "Convolution/NumberTheoreticTransform.hpp"
#include <cstdint>

/**
 * @brief https://tkmst201.github.io/Library/Convolution/NumberTheoreticTransform.hpp
 */
template<int MOD, int PRIMITIVE_ROOT>
struct NumberTheoreticTransform {
	static_assert(MOD > 0);
	static_assert(PRIMITIVE_ROOT > 0);
	
private:
	using uint32 = std::uint32_t;
	using calc_type = long long;
	
public:
	template<typename T>
	static std::vector<T> multiply(const std::vector<T> & a, const std::vector<T> & b) {
		static_assert(std::is_integral<T>::value);
		if (a.empty() || b.empty()) return {};
		const uint32 n_ = a.size() + b.size() - 1;
		uint32 n = 1;
		while (n < n_) n <<= 1;
		{
			uint32 two_exp = 0;
			int tm = MOD - 1;
			while (tm > 0 && (~tm & 1)) ++two_exp, tm >>= 1;
			assert((1u << two_exp) >= n);
		}
		std::vector<T> c(n, 0);
		for (uint32 i = 0; i < a.size(); ++i) c[i] = a[i] % MOD + (a[i] >= 0 ? 0 : MOD);
		ntt(c);
		std::vector<T> d(n, 0);
		for (uint32 i = 0; i < b.size(); ++i) d[i] = b[i] % MOD + (b[i] >= 0 ? 0 : MOD);
		ntt(d);
		const int ninv = tk::mod_inv<int>(n, MOD);
		for (uint32 i = 0; i < n; ++i) c[i] = static_cast<calc_type>(c[i]) * d[i] % MOD * ninv % MOD;
		d.clear();
		ntt(c, true);
		c.resize(a.size() + b.size() - 1);
		return c;
	}
	
private:
	template<typename T>
	static void ntt(std::vector<T> & a, bool inv = false) {
		const uint32 n = a.size();
		int nroot = tk::mod_pow<calc_type>(PRIMITIVE_ROOT, (MOD - 1) / n, MOD);
		if (inv) nroot = tk::mod_inv(nroot, MOD);
		for (uint32 w = n; w > 1; w >>= 1) {
			const uint32 m = w >> 1;
			std::vector<int> omega(m, 0);
			omega[0] = 1;
			for (uint32 i = 1; i < m; ++i) omega[i] = static_cast<calc_type>(omega[i - 1]) * nroot % MOD;
			const int half = static_cast<calc_type>(omega.back()) * nroot % MOD;
			for (uint32 p = 0; p < n; p += w) {
				for (uint32 i = p; i < p + m; ++i) {
					const calc_type x = a[i], y = a[i + m];
					a[i] = (x + y) % MOD;
					a[i + m] = (x + y * half % MOD) % MOD * omega[i - p] % MOD;
				}
			}
			nroot = static_cast<calc_type>(nroot) * nroot % MOD;
		}
		bit_reverse(a);
	}
	
	template<typename T>
	static void bit_reverse(std::vector<T> & a) noexcept {
		const uint32 n = a.size();
		for (uint32 i = 1, j = 0; i < n - 1; ++i) {
			for (uint32 k = n >> 1; k > (j ^= k); k >>= 1);
			if (i < j) std::swap(a[i], a[j]);
		}
	}
};


#line 1 "Mathematics/garner.hpp"



#line 7 "Mathematics/garner.hpp"

#line 1 "Mathematics/euclid.hpp"



#line 6 "Mathematics/euclid.hpp"
#include <tuple>
#line 8 "Mathematics/euclid.hpp"
#include <cmath>

/**
 * @brief https://tkmst201.github.io/Library/Mathematics/euclid.hpp
 */
namespace tk {
template<typename T>
constexpr T gcd(T a, T b) noexcept {
	static_assert(std::is_integral<T>::value);
	assert(a >= 0);
	assert(b >= 0);
	while (b != 0) {
		const T t = a % b;
		a = b; b = t;
	}
	return a;
}

template<typename T>
constexpr T lcm(T a, T b) noexcept {
	static_assert(std::is_integral<T>::value);
	assert(a >= 0);
	assert(b >= 0);
	if (a == 0 || b == 0) return 0;
	return a / gcd(a, b) * b;
}

template<typename T>
constexpr std::tuple<T, T, T> ext_gcd(T a, T b) noexcept {
	static_assert(std::is_integral<T>::value);
	static_assert(std::is_signed<T>::value);
	assert(a != 0);
	assert(b != 0);
	T a1 = (a > 0) * 2 - 1, a2 = 0, b1 = 0, b2 = (b > 0) * 2 - 1;
	a = std::abs(a);
	b = std::abs(b);
	while (b > 0) {
		const T q = a / b;
		T tmp = a - q * b; a = b; b = tmp;
		tmp = a1 - q * b1; a1 = b1; b1 = tmp;
		tmp = a2 - q * b2; a2 = b2; b2 = tmp;
	}
	return {a, a1, a2};
}
} // namespace tk



#line 10 "Mathematics/garner.hpp"

/**
 * @brief https://tkmst201.github.io/Library/Mathematics/garner.hpp
 */
namespace tk {
template<typename T>
bool pre_garner(std::vector<T> & b, std::vector<T> & m) noexcept {
	static_assert(std::is_integral<T>::value);
	static_assert(std::is_signed<T>::value);
	for (int i = 0; i < static_cast<int>(b.size()); ++i) {
		b[i] = b[i] % m[i] + (b[i] >= 0 ? 0 : m[i]);
		for (int j = 0; j < i; ++j) {
			T g = gcd(m[i], m[j]);
			if ((b[i] - b[j]) % g != 0) return false;
			m[i] /= g; m[j] /= g;
			T gi = gcd(g, m[i]), gj = g / gi;
			do {
				g = gcd(gi, gj);
				gi *= g; gj /= g;
			} while (g != 1);
			m[i] *= gi; m[j] *= gj;
			b[i] %= m[i]; b[j] %= m[j];
		}
	}
	return true;
}

template<typename T, typename U>
T garner(const std::vector<T> & b, const std::vector<T> & m, const T M) {
	static_assert(std::is_integral<T>::value);
	assert(b.size() == m.size());
	const int n = b.size();
	assert(n > 0);
	{
		T g = 0;
		for (auto v : m) {
			assert(v > 0);
			g = gcd(g, v);
		}
		assert(n == 1 || g == 1);
	}
	assert(M > 0);
	std::vector<T> sum(n + 1, 0), ip(n + 1, 1);
	for (int i = 0; i < n; ++i) {
		if (m[i] == 1) continue;
		U t = (b[i] % m[i] + static_cast<U>(2) * m[i] - sum[i]) % m[i] * mod_inv(ip[i], m[i]) % m[i];
		for (int j = i + 1; j < n; ++j) {
			sum[j] = (sum[j] + ip[j] * t) % m[j];
			ip[j] = static_cast<U>(ip[j]) * m[i] % m[j];
		}
		sum[n] = (sum[n] + ip[n] * t % M) % M;
		ip[n] = static_cast<U>(ip[n]) * m[i] % M; 
	}
	return sum.back();
}
} // namespace tk



#line 6 "Convolution/NumberTheoreticTransform_AnyMod.hpp"

#line 11 "Convolution/NumberTheoreticTransform_AnyMod.hpp"

/**
 * @brief https://tkmst201.github.io/Library/Convolution/NumberTheoreticTransform_AnyMod.hpp
 */
template<int MOD>
struct NumberTheoreticTransform_AnyMod {
	static_assert(MOD > 0);
	
private:
	using calc_type = long long;
	using uint32 = std::uint32_t;
	
public:
	template<typename T>
	static std::vector<T> multiply(const std::vector<T> & a, const std::vector<T> & b) {
		static_assert(std::is_integral<T>::value);
		static_assert(std::is_signed<T>::value);
		for (uint32 i = 0; i < a.size(); ++i) assert(a[i] >= 0);
		for (uint32 i = 0; i < b.size(); ++i) assert(b[i] >= 0);
		std::vector<T> m;
		auto ntt1_res = NumberTheoreticTransform<1'224'736'769, 3>::multiply(a, b);
		m.emplace_back(1'224'736'769);
		auto ntt2_res = NumberTheoreticTransform<469'762'049, 3>::multiply(a, b);
		m.emplace_back(469'762'049);
		auto ntt3_res = NumberTheoreticTransform<167'772'161, 3>::multiply(a, b);
		m.emplace_back(167'772'161);
		// auto ntt4_res = NumberTheoreticTransform<998'244'353, 3>::multiply(a, b);

		// m.emplace_back(998'244'353);

		
		std::vector<T> c(m.size());
		std::vector<T> res(ntt1_res.size());
		for (uint32 i = 0; i < res.size(); ++i) {
			c[0] = ntt1_res[i];
			c[1] = ntt2_res[i];
			c[2] = ntt3_res[i];
			// c[3] = ntt4_res[i];

			res[i] = tk::garner<T, calc_type>(c, m, MOD);
		}
		return res;
	}
};

任意 mod 数論変換 (任意 mod NTT) (Convolution/NumberTheoreticTransform_AnyMod.hpp)

概要