Fast Fourier Transform
(convolution/fft.hpp)

• View this file on GitHub
• Last update: 2024-01-08 00:27:17+09:00
• Include: #include "convolution/fft.hpp"

Description

高速 Fourier 変換 (FFT) は数列の離散 Fourier 変換を計算するアルゴリズムであり，2つの数列の畳み込みを高速に計算するのに用いられる．

この実装では Cooley-Tukey のアルゴリズムを用いている．

Operations

• void fft(vector<complex<double>> a)
  • 数列 $a$ を高速 Fourier 変換する
  • 時間計算量: $O(n\log n)$
• void ifft(vector<complex<double>> a)
  • 数列 $a$ を逆高速 Fourier 変換する
  • 時間計算量: $O(n\log n)$
• vector<double> convolution(vector<T> a, vector<T> b)
  • 数列 $a$ と $b$ の畳み込みを計算する
  • 時間計算量: $O(n\log n)$

Verified with

• test/yosupo/frequency_table_of_tree_distance.test.cpp

Code

#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;
}

Back to top page