test/convolution_fft.test.cpp

test/convolution_fft.test.cpp · GitHub で見る

依存ライブラリ

コード

#define PROBLEM "https://yukicoder.me/problems/no/1307"
#include "../convolution/convolution_fft.hpp"
#include <cmath>
#include <iostream>
#include <vector>

int main() {
	std::cin.tie(nullptr)->sync_with_stdio(false);
	int n, q;
	std::cin >> n >> q;
	std::vector<double> a(n), b(n);
	for (auto &x : a) std::cin >> x;
	while (q--) {
		int r;
		std::cin >> r;
		++b[(n - r) % n];
	}
	auto c = convolution_fft(a, b);
	for (int i = 0; i < n; ++i)
		std::cout << std::llround(c[i] + (i + n < int(c.size()) ? c[i + n] : 0)) << (i + 1 == n ? '\n' : ' ');
}

#line 1 "test/convolution_fft.test.cpp"
#define PROBLEM "https://yukicoder.me/problems/no/1307"
#line 1 "convolution/convolution_fft.hpp"



#line 1 "components/fft.hpp"



#include <algorithm>
#include <bit>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstddef>
#include <numbers>
#include <vector>

template <class Real>
void fft(std::vector<std::complex<Real>> &a, bool inverse = false) {
	std::size_t n = a.size();
	assert(!n || std::has_single_bit(n));
	for (std::size_t i = 1, j = 0; i < n; ++i) {
		std::size_t bit = n >> 1;
		for (; j & bit; bit >>= 1) j ^= bit;
		j ^= bit;
		if (i < j) std::swap(a[i], a[j]);
	}
	for (std::size_t width = 2; width <= n; width *= 2) {
		Real angle = (inverse ? -2 : 2) * std::numbers::pi_v<Real> / Real(width);
		std::complex<Real> step(std::cos(angle), std::sin(angle));
		for (std::size_t start = 0; start < n; start += width) {
			std::complex<Real> w(1);
			for (std::size_t j = 0; j < width / 2; ++j, w *= step) {
				auto x = a[start + j], y = a[start + j + width / 2] * w;
				a[start + j] = x + y;
				a[start + j + width / 2] = x - y;
			}
		}
	}
	if (inverse)
		for (auto &x : a) x /= Real(n);
}


#line 9 "convolution/convolution_fft.hpp"
#include <type_traits>
#line 11 "convolution/convolution_fft.hpp"

template <class Real>
std::vector<std::complex<Real>> convolution_fft(const std::vector<std::complex<Real>> &a, const std::vector<std::complex<Real>> &b) {
	if (a.empty() || b.empty()) return {};
	std::size_t length = a.size() + b.size() - 1, n = std::bit_ceil(length);
	std::vector<std::complex<Real>> x(a.begin(), a.end()), y(b.begin(), b.end());
	x.resize(n);
	y.resize(n);
	fft(x, false);
	if (&a == &b) y = x;
	else fft(y, false);
	for (std::size_t i = 0; i < n; ++i) x[i] *= y[i];
	fft(x, true);
	x.resize(length);
	return x;
}
template <class Real>
	requires std::is_floating_point_v<Real>
std::vector<Real> convolution_fft(const std::vector<Real> &a, const std::vector<Real> &b) {
	std::vector<std::complex<Real>> x(a.begin(), a.end()), y(b.begin(), b.end());
	auto values = &a == &b ? convolution_fft(x, x) : convolution_fft(x, y);
	std::vector<Real> result;
	result.reserve(values.size());
	for (const auto &x : values) result.push_back(x.real());
	return result;
}


#line 4 "test/convolution_fft.test.cpp"
#include <iostream>
#line 6 "test/convolution_fft.test.cpp"

int main() {
	std::cin.tie(nullptr)->sync_with_stdio(false);
	int n, q;
	std::cin >> n >> q;
	std::vector<double> a(n), b(n);
	for (auto &x : a) std::cin >> x;
	while (q--) {
		int r;
		std::cin >> r;
		++b[(n - r) % n];
	}
	auto c = convolution_fft(a, b);
	for (int i = 0; i < n; ++i)
		std::cout << std::llround(c[i] + (i + n < int(c.size()) ? c[i + n] : 0)) << (i + 1 == n ? '\n' : ' ');
}