数論変換 (NTT)
(Convolution/NumberTheoreticTransform.hpp)

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

概要

特殊な整数 $M$ と $2$ つの整数係数多項式 $f(x), g(x)$ に対して、積 $f(x)g(x) \bmod{M}$ を計算します。計算量は $\Theta(N\log{N})$ (積の次数を $N$ ) です。
基数 $2$ の時間間引き Cooley-Tukey 型アルゴリズムを用いています。

静的メンバ関数

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

制約

$M > 0$
適当な整数 $p$ を用いて $M - 1 = Np$ ($N$ は $|a| + |b| - 1$ 以上の最小の $2$ 冪) と表すことが可能
PRIMITIVE_ROOT は $M$ の原始根
T は int, unsigned int, long long int, unsigned long long

計算量

$N$ を $|a| + |b| - 1$ 以上の最小の $2$ 冪として $\Theta(N\log{N})$

使用例

代表的な MOD とその原始根の組み合わせの例を下に載せておきます。

NumberTheoreticTransform<998'244'353, 3> // 2^23 | (mod - 1)
NumberTheoreticTransform<1'224'736'769, 3> // 2^24 | (mod - 1)
NumberTheoreticTransform<469'762'049, 3> // 2^26 | (mod - 1)
NumberTheoreticTransform<167'772'161, 3> // 2^25 | (mod - 1)

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

int main() {
	vector<int> a({0, 1, 2, 3}), b({2, 3, 4});
	auto c = NumberTheoreticTransform<998244353, 3>::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 16 17 12
	// 0 0 0
	//   2 3 4
	//     4 6 8
	//       6 9 12
	// ==============
	// 0 2 7 16 17 12
	
	vector<int> a2({4}), b2({3, 2});
	auto c2 = NumberTheoreticTransform<167772161, 3>::multiply(a2, b2);
	cout << "size = " << c2.size() << endl; // 2
	for (int i = 0; i < c2.size(); ++i) cout << c2[i] << " "; cout << endl; // 12 8
	// 12 
	//    8
	// ==============
	// 12 8
	
	vector<int> a3({-3, 1}), b3({-2, -9});
	auto c3 = NumberTheoreticTransform<5, 3>::multiply(a3, b3);
	cout << "size = " << c3.size() << endl; // 3
	for (int i = 0; i < c3.size(); ++i) cout << c3[i] << " "; cout << endl; // 1 0 1
	// 1(6) 3(-2)
	//      2(27) 1(-9)
	// ==============
	// 1    0     1
	
	vector<int> a4, b4({1, 2, 3});
	auto c4 = NumberTheoreticTransform<469'762'049, 3>::multiply(a4, b4);
	cout << "size = " << c4.size() << endl; // 0
}

参考

2020/07/27: https://math314.hateblo.jp/entry/2015/05/07/014908
2020/07/27: https://ja.wikibooks.org/wiki/%E5%88%9D%E7%AD%89%E6%95%B4%E6%95%B0%E8%AB%96/%E5%8E%9F%E5%A7%8B%E6%A0%B9%E3%81%A8%E6%8C%87%E6%95%B0
2020/07/27: http://wwwa.pikara.ne.jp/okojisan/stockham/cooley-tukey.html

Depends on

$\mod{m}$ での冪と逆元 (Mathematics/mod_pow_inv.hpp)

Required by

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

Verified with

Code

#ifndef INCLUDE_GUARD_NUMBER_THEORETIC_TRANSFORM_HPP
#define INCLUDE_GUARD_NUMBER_THEORETIC_TRANSFORM_HPP

#include "Mathematics/mod_pow_inv.hpp"

#include <vector>
#include <utility>
#include <cassert>
#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]);
		}
	}
};

#endif // INCLUDE_GUARD_NUMBER_THEORETIC_TRANSFORM_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]);
		}
	}
};

数論変換 (NTT) (Convolution/NumberTheoreticTransform.hpp)

概要