feat: Add Fast Fourier Transform (FFT) implementation and convolution functions

Author: weilycoder

test/multiplication_of_big_integers.fft.test.cpp

@@ -0,0 +1,62 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/multiplication_of_big_integers"
+
+#include "../weilycoder/poly/fft_convolve.hpp"
+#include <algorithm>
+#include <iostream>
+#include <string>
+#include <vector>
+using namespace std;
+using namespace weilycoder;
+
+static void solve() {
+  string a, b;
+  cin >> a >> b;
+
+  reverse(a.begin(), a.end()), reverse(b.begin(), b.end());
+
+  bool negative = false;
+  if (a.back() == '-')
+    negative = !negative, a.pop_back();
+  if (b.back() == '-')
+    negative = !negative, b.pop_back();
+
+  vector<double> A, B;
+  for (char c : a)
+    A.push_back(c - '0');
+  for (char c : b)
+    B.push_back(c - '0');
+
+  vector<double> C = fft_convolve_real(A, B);
+  vector<int> result;
+  int carry = 0;
+  for (double x : C) {
+    int val = static_cast<int>(x + 0.5) + carry;
+    result.push_back(val % 10);
+    carry = val / 10;
+  }
+  while (carry) {
+    result.push_back(carry % 10);
+    carry /= 10;
+  }
+  while (!result.empty() && result.back() == 0)
+    result.pop_back();
+  if (result.empty()) {
+    cout << 0 << '\n';
+  } else {
+    if (negative)
+      cout << '-';
+    for (auto it = result.rbegin(); it != result.rend(); ++it)
+      cout << *it;
+    cout << '\n';
+  }
+}
+
+int main() {
+  cin.tie(nullptr)->sync_with_stdio(false);
+  cin.exceptions(cin.failbit | cin.badbit);
+  size_t t;
+  cin >> t;
+  while (t--)
+    solve();
+  return 0;
+}

weilycoder/poly/fft.hpp

@@ -0,0 +1,43 @@
+#ifndef WEILYCODER_POLY_FFT_HPP
+#define WEILYCODER_POLY_FFT_HPP
+
+#include "fft_utility.hpp"
+#include <complex>
+#include <cstddef>
+#include <vector>
+
+/**
+ * @file fft.hpp
+ * @brief Implementation of Fast Fourier Transform (FFT) and its inverse.
+ */
+
+namespace weilycoder {
+/**
+ * @brief In-place iterative Cooley–Tukey FFT / inverse FFT on a complex vector.
+ *        The length of the input vector should be a power of two.
+ * @tparam on Direction of the transform: 1 for FFT, -1 for inverse FFT.
+ * @tparam float_t Floating-point type for computations (default: double).
+ * @param y Vector of complex numbers to be transformed in-place.
+ */
+template <int32_t on = 1, typename float_t = double>
+void fft(std::vector<std::complex<float_t>> &y) {
+  static_assert(on == 1 || on == -1, "on must be 1 or -1");
+  fft_change(y);
+  for (size_t h = 2; h <= y.size(); h <<= 1) {
+    std::complex<float_t> wn(cos(2 * PI<float_t> / h), sin(on * 2 * PI<float_t> / h));
+    for (size_t j = 0; j < y.size(); j += h) {
+      std::complex<float_t> w(1, 0);
+      for (size_t k = j; k < j + (h >> 1); ++k, w *= wn) {
+        std::complex<float_t> u = y[k], t = w * y[k + (h >> 1)];
+        y[k] = u + t, y[k + (h >> 1)] = u - t;
+      }
+    }
+  }
+  size_t len = y.size();
+  if constexpr (on == -1)
+    for (size_t i = 0; i < len; ++i)
+      y[i] /= len;
+}
+} // namespace weilycoder
+
+#endif

weilycoder/poly/fft_convolve.hpp

@@ -0,0 +1,66 @@
+#ifndef WEILYCODER_POLY_FFT_CONVOLVE_HPP
+#define WEILYCODER_POLY_FFT_CONVOLVE_HPP
+
+#include "fft.hpp"
+#include <complex>
+#include <cstddef>
+#include <vector>
+
+/**
+ * @file fft_convolve.hpp
+ * @brief Functions for performing convolution using Fast Fourier Transform (FFT).  
+ */
+
+namespace weilycoder {
+/**
+ * @brief Perform convolution of two complex vectors using FFT.
+ * @tparam float_t Floating-point type for computations (default: double).
+ * @param a First input vector of complex numbers.
+ * @param b Second input vector of complex numbers.
+ * @return Vector containing the convolution result.
+ */
+template <typename float_t = double>
+std::vector<std::complex<float_t>> fft_convolve(std::vector<std::complex<float_t>> a,
+                                                std::vector<std::complex<float_t>> b) {
+  size_t n = 1;
+  while (n < a.size() + b.size() - 1)
+    n <<= 1;
+  a.resize(n), b.resize(n);
+  fft(a), fft(b);
+  for (size_t i = 0; i < n; ++i)
+    a[i] *= b[i];
+  return fft<-1>(a), a;
+}
+
+/**
+ * @brief Perform convolution of two real-valued vectors using FFT.
+ * @tparam float_t Floating-point type for computations (default: double).
+ * @param a First input vector of real numbers.
+ * @param b Second input vector of real numbers.
+ * @return Vector containing the convolution result.
+ * @note This function uses a technique that combines two real sequences into
+ *       a single complex sequence to perform the convolution more efficiently.
+ */
+template <typename float_t = double>
+std::vector<float_t> fft_convolve_real(const std::vector<float_t> &a,
+                                       const std::vector<float_t> &b) {
+  size_t n = 1;
+  while (n < a.size() + b.size() - 1)
+    n <<= 1;
+  std::vector<std::complex<float_t>> F(n);
+  for (size_t i = 0; i < a.size(); ++i)
+    F[i].real(a[i]);
+  for (size_t i = 0; i < b.size(); ++i)
+    F[i].imag(b[i]);
+  fft(F);
+  for (size_t i = 0; i < n; ++i)
+    F[i] *= F[i];
+  fft<-1>(F);
+  std::vector<float_t> result(a.size() + b.size() - 1);
+  for (size_t i = 0; i < result.size(); ++i)
+    result[i] = F[i].imag() / 2;
+  return result;
+}
+} // namespace weilycoder
+
+#endif

weilycoder/poly/fft_utility.hpp

@@ -0,0 +1,46 @@
+#ifndef WEILYCODER_POLY_FFT_UTILITY_HPP
+#define WEILYCODER_POLY_FFT_UTILITY_HPP
+
+#include <complex>
+#include <cstddef>
+#include <vector>
+
+/**
+ * @file fft_utility.hpp
+ * @brief Utility functions and constants for Fast Fourier Transform (FFT)
+ */
+
+namespace weilycoder {
+/**
+ * @brief Alias for the commonly used complex number type with double precision.
+ */
+using complex_t = std::complex<double>;
+
+/**
+ * @brief Compile-time constant for π (pi) parameterized by numeric type.
+ * @tparam T Numeric type (e.g., float, double, long double).
+ */
+template <typename T> constexpr T PI = T(3.1415926535897932384626433832795l);
+
+/**
+ * @brief Perform in-place bit-reversal permutation (index reordering) required by
+ *        iterative FFT.
+ * @tparam T Element type stored in the vector. Must be Swappable and efficiently
+ *         copyable.
+ * @param a Vector to be permuted in-place. Its size should typically be a power
+ *          of two for use with the FFT implementation in this file.
+ */
+template <typename T> void fft_change(std::vector<T> &a) {
+  size_t n = a.size();
+  std::vector<size_t> rev(n);
+  for (size_t i = 0; i < n; ++i) {
+    rev[i] = rev[i >> 1] >> 1;
+    if (i & 1)
+      rev[i] |= n >> 1;
+    if (i < rev[i])
+      swap(a[i], a[rev[i]]);
+  }
+}
+} // namespace weilycoder
+
+#endif