fft.hpp

#pragma once
 
#include <cmath>
#include <complex>
#include <iterator>
#include <vector>
 
#include "cplib/num/mmint.hpp"
#include "cplib/num/pow.hpp"
#include "cplib/port/bit.hpp"
 
namespace cplib {
 
namespace impl {
 
template <typename T>
std::vector<T> twiddling_factors(const std::vector<T>& roots) {
  std::size_t n = 1 << (roots.size() - 2);
  std::vector<T> w, w_stack;
  w.reserve(n);
  w_stack.reserve(roots.size());
  w.push_back(roots.front());
  w_stack.push_back(roots.front());
  // For any give i, the k-th element in the stack is i with only the highest k 1-bits.
  // Each element is the product of its lowest bit and other bits, and thus is the product of at most log2(N) roots.
  // Thus for complex numbers the error in twiddling factors is at most log2(N) eps.
  for (std::size_t i = 1; i < n; i++) {
    w_stack.resize(port::popcount(i));
    w_stack.push_back(w_stack.back() * roots[roots.size() - 1 - port::countr_zero(i)]);
    w.push_back(w_stack.back());
  }
  return w;
}
 
}  // namespace impl
 
template <typename T>
struct radix2_fft_root {};
 
template <>
struct radix2_fft_root<MMInt<998244353>> {
  using mint = MMInt<998244353>;
  static mint get(int n) { return pow(mint(3), 119 << (23 - n)); }
};
 
template <>
struct radix2_fft_root<MMInt64<4512606826625236993>> {
  using mint = MMInt64<4512606826625236993>;
  static mint get(int n) { return pow(mint(7), 501ull << (53 - n)); }
};
 
template <typename Float>
struct radix2_fft_root<std::complex<Float>> {
  constexpr static std::complex<Float> get(size_t n) {
    constexpr long double tau = atanl(1) * 8;  // In C++20, std::pi_v<Float> * 2
    return std::polar<Float>(1, tau / (1ull << n));
  }
};
 
template <typename RandomIt>
void fft_inplace(RandomIt first, RandomIt last) {
  using usize = std::size_t;
  const usize n = std::distance(first, last);
  assert(port::has_single_bit(n));
  int log2n = port::countr_zero(n);
  using T = typename std::iterator_traits<RandomIt>::value_type;
  std::vector<T> roots(log2n + 1);
  for (int i = 0; i <= log2n; i++) {
    roots[i] = radix2_fft_root<T>::get(i);
  }
  for (int stage = log2n - 1; stage >= 0; stage--) {
    usize len = usize(1) << stage;
    std::vector<T> w = impl::twiddling_factors(roots);
    for (usize block = 0; block < n; block += len * 2) {
      for (usize offset = 0; offset < len; offset++) {
        usize i = block + offset, j = i + len;
        T tmp = (first[i] - first[j]) * w[offset];
        first[i] += first[j];
        first[j] = tmp;
      }
    }
    roots.pop_back();
  }
}
 
template <typename RandomIt>
void ifft_inplace(RandomIt first, RandomIt last) {
  using usize = std::size_t;
  const usize n = std::distance(first, last);
  assert(port::has_single_bit(n));
  int log2n = port::countr_zero(n);
  using T = typename std::iterator_traits<RandomIt>::value_type;
  T one = radix2_fft_root<T>::get(0);
  std::vector<T> roots{one};
  for (int stage = 0; stage < log2n; stage++) {
    usize len = usize(1) << stage;
    roots.push_back(one / radix2_fft_root<T>::get(stage + 1));
    std::vector<T> w = impl::twiddling_factors(roots);
    for (usize block = 0; block < n; block += len * 2) {
      for (usize offset = 0; offset < len; offset++) {
        usize i = block + offset, j = i + len;
        first[j] *= w[offset];
        T tmp = first[i] - first[j];
        first[i] += first[j];
        first[j] = tmp;
      }
    }
  }
  T half = one / (one + one);
  T n_inv = pow(half, log2n);
  for (usize i = 0; i < n; i++) {
    first[i] *= n_inv;
  }
}
 
template <typename T>
void fft_inplace(std::vector<T>& a) {
  fft_inplace(a.begin(), a.end());
}
 
template <typename T>
void ifft_inplace(std::vector<T>& a) {
  ifft_inplace(a.begin(), a.end());
}
 
}  // namespace cplib