sotanishy's code snippets for competitive programming
#include "convolution/fft.hpp"
高速 Fourier 変換 (FFT) は数列の離散 Fourier 変換を計算するアルゴリズムであり,2つの数列の畳み込みを高速に計算するのに用いられる.
この実装では Cooley-Tukey のアルゴリズムを用いている.
void fft(vector<complex<double>> a)
void ifft(vector<complex<double>> a)
vector<double> convolution(vector<T> a, vector<T> b)
#pragma once
#include <bit>
#include <complex>
#include <numbers>
#include <vector>
void fft(std::vector<std::complex<double>>& a) {
const int n = a.size();
for (int m = n; m > 1; m >>= 1) {
double ang = 2.0 * std::numbers::pi / m;
std::complex<double> omega(cos(ang), sin(ang));
for (int s = 0; s < n / m; ++s) {
std::complex<double> w(1, 0);
for (int i = 0; i < m / 2; ++i) {
auto l = a[s * m + i];
auto r = a[s * m + i + m / 2];
a[s * m + i] = l + r;
a[s * m + i + m / 2] = (l - r) * w;
w *= omega;
}
}
}
}
void ifft(std::vector<std::complex<double>>& a) {
const int n = a.size();
for (int m = 2; m <= n; m <<= 1) {
double ang = -2.0 * std::numbers::pi / m;
std::complex<double> omega(cos(ang), sin(ang));
for (int s = 0; s < n / m; ++s) {
std::complex<double> w(1, 0);
for (int i = 0; i < m / 2; ++i) {
auto l = a[s * m + i];
auto r = a[s * m + i + m / 2] * w;
a[s * m + i] = l + r;
a[s * m + i + m / 2] = l - r;
w *= omega;
}
}
}
}
template <typename T>
std::vector<double> convolution(const std::vector<T>& a,
const std::vector<T>& b) {
const int size = a.size() + b.size() - 1;
const int n = std::bit_ceil((unsigned int)size);
std::vector<std::complex<double>> na(a.begin(), a.end()),
nb(b.begin(), b.end());
na.resize(n);
nb.resize(n);
fft(na);
fft(nb);
for (int i = 0; i < n; ++i) na[i] *= nb[i];
ifft(na);
std::vector<double> ret(size);
for (int i = 0; i < size; ++i) ret[i] = na[i].real() / n;
return ret;
}
#line 2 "convolution/fft.hpp"
#include <bit>
#include <complex>
#include <numbers>
#include <vector>
void fft(std::vector<std::complex<double>>& a) {
const int n = a.size();
for (int m = n; m > 1; m >>= 1) {
double ang = 2.0 * std::numbers::pi / m;
std::complex<double> omega(cos(ang), sin(ang));
for (int s = 0; s < n / m; ++s) {
std::complex<double> w(1, 0);
for (int i = 0; i < m / 2; ++i) {
auto l = a[s * m + i];
auto r = a[s * m + i + m / 2];
a[s * m + i] = l + r;
a[s * m + i + m / 2] = (l - r) * w;
w *= omega;
}
}
}
}
void ifft(std::vector<std::complex<double>>& a) {
const int n = a.size();
for (int m = 2; m <= n; m <<= 1) {
double ang = -2.0 * std::numbers::pi / m;
std::complex<double> omega(cos(ang), sin(ang));
for (int s = 0; s < n / m; ++s) {
std::complex<double> w(1, 0);
for (int i = 0; i < m / 2; ++i) {
auto l = a[s * m + i];
auto r = a[s * m + i + m / 2] * w;
a[s * m + i] = l + r;
a[s * m + i + m / 2] = l - r;
w *= omega;
}
}
}
}
template <typename T>
std::vector<double> convolution(const std::vector<T>& a,
const std::vector<T>& b) {
const int size = a.size() + b.size() - 1;
const int n = std::bit_ceil((unsigned int)size);
std::vector<std::complex<double>> na(a.begin(), a.end()),
nb(b.begin(), b.end());
na.resize(n);
nb.resize(n);
fft(na);
fft(nb);
for (int i = 0; i < n; ++i) na[i] *= nb[i];
ifft(na);
std::vector<double> ret(size);
for (int i = 0; i < size; ++i) ret[i] = na[i].real() / n;
return ret;
}