Fast Fourier Transform
(src/algebra/fft.hpp)

View this file on GitHub
Last update: 2022-04-06 15:02:09+09:00
Include: #include "src/algebra/fft.hpp"

Depends on

Required by

Verified with

Code

#pragma once

/**
 * @file fft.hpp
 * @brief Fast Fourier Transform
 */

#include <vector>

#include "complex.hpp"
#include "lib/cxx17"
#include "src/number_theory/ext_gcd.hpp"
#include "src/number_theory/primitive_root.hpp"
#include "src/utils/sfinae.hpp"

namespace workspace {

namespace _fft_impl {

template <class _Tp, bool = std::is_floating_point<_Tp>::value, class = void>
struct to_float {
  using type = double;
};

template <class _Tp> struct to_float<_Tp, true> { using type = _Tp; };

// template <class _Tp>
// struct to_float<_Tp, false, std::enable_if_t<sizeof(_Tp) <= sizeof(float)>> {
//   using type = float;
// };

template <class _Tp>
struct to_float<_Tp, false, std::enable_if_t<(sizeof(_Tp) > sizeof(double))>> {
  using type = long double;
};

// Assume ntt-friendly mod.
template <class _Tp> struct field {
  using type = std::conditional_t<has_mod<_Tp>::value, _Tp,
                                  complex<typename to_float<_Tp>::type>>;
};

template <class _Tp> struct field<complex<_Tp>> : field<_Tp> {};

// Modular
template <class _Tp, int _Nm = 29, bool = has_mod<_Tp>::value> struct coef {
  _Tp s[_Nm], is[_Nm], ip2[_Nm];

  _CXX14_CONSTEXPR coef() : s{}, is{}, ip2{1, (1 + _Tp::mod) / 2} {
    if (_Tp::mod < 2) return;

    int cnt2 = std::min(__builtin_ctz(_Tp::mod - 1), _Nm + 1);

    _Tp e = 1;
    _Tp w = primitive_root(_Tp::mod);
    for (auto p = (_Tp::mod - 1) >> cnt2; p; p >>= 1, w *= w)
      if (p & 1) e *= w;

    _Tp ie = ext_gcd(decltype(_Tp::mod)(e), _Tp::mod).first;
    _Tp es[_Nm]{}, ies[_Nm]{};

    for (int i = cnt2; i >= 2; i--) {
      es[i - 2] = e, e *= e;
      ies[i - 2] = ie, ie *= ie;
    }

    e = ie = 1;
    for (int i = 0; i < cnt2 - 1; i++) {
      s[i] = es[i] * e, e *= ies[i];
      is[i] = ies[i] * ie, ie *= es[i];
    }

    for (int i = 1; i < _Nm - 1; ++i) ip2[i + 1] = ip2[i] * ip2[1];
  }
};

// Complex
template <class _Tp, int _Nm> struct coef<_Tp, _Nm, false> {
  _Tp s[_Nm], is[_Nm], ip2[_Nm];

  static_assert(_Nm < 30);

  _CXX14_CONSTEXPR static _Tp es[29] = {
      {0, 1},
      {0.70710678118654752438189403651, 0.70710678118654752443610414514},
      {0.92387953251128675610142140795, 0.38268343236508977172325753068},
      {0.98078528040323044911909938781, 0.19509032201612826785692544201},
      {0.99518472667219688623102546998, 0.09801714032956060199569840382},
      {0.99879545620517239270077028412, 0.04906767432741801425693899119},
      {0.99969881869620422009748220149, 0.02454122852291228803212346128},
      {0.99992470183914454093764001552, 0.01227153828571992607945510345},
      {0.99998117528260114264494415325, 0.00613588464915447535972750246},
      {0.99999529380957617150137498041, 0.00306795676296597627029751672},
      {0.99999882345170190993313003025, 0.00153398018628476561237225788},
      {0.99999970586288221914474799723, 0.00076699031874270452695124765},
      {0.99999992646571785113833452651, 0.00038349518757139558906815188},
      {0.99999998161642929381167504976, 0.00019174759731070330743679009},
      {0.99999999540410731290905263501, 0.00009587379909597734587360460},
      {0.99999999885102682753608427379, 0.00004793689960306688454884772},
      {0.99999999971275670682981095982, 0.00002396844980841821872882467},
      {0.99999999992818917670745273995, 0.00001198422490506970642183282},
      {0.99999999998204729416331065783, 0.00000599211245264242784278378},
      {0.99999999999551182356793271877, 0.00000299605622633466075058210},
      {0.99999999999887795586487812538, 0.00000149802811316901122883643},
      {0.99999999999971948897977205850, 0.00000074901405658471572113723},
      {0.99999999999992987223139048746, 0.00000037450702829238412391495},
      {0.99999999999998246807140014902, 0.00000018725351414619534486931},
      {0.99999999999999561700429751010, 0.00000009362675707309808280024},
      {0.99999999999999890425107437752, 0.00000004681337853654909269501},
      {0.99999999999999972607632112153, 0.00000002340668926827455275977},
      {0.99999999999999993153263280754, 0.00000001170334463413727718121},
      {0.99999999999999998286960567472, 0.00000000585167231706863869077}};

  _CXX14_CONSTEXPR coef() : s{}, is{}, ip2{1, .5} {
    _Tp ies[_Nm];
    for (int i = 0; i < _Nm; ++i) ies[i] = _Tp(1) / es[i];

    _Tp e = 1, ie = 1;
    for (int i = 0; i < _Nm; i++) {
      s[i] = es[i] * e, e *= ies[i];
      is[i] = ies[i] * ie, ie *= es[i];
    }

    for (int i = 1; i < _Nm - 1; ++i) ip2[i + 1] = ip2[i] * ip2[1];
  }
};

}  // namespace _fft_impl

template <bool _Inverse = false, class _Iterator>
void fft(_Iterator __first, _Iterator __last) noexcept {
  using value_type = typename std::iterator_traits<_Iterator>::value_type;
  using difference_type =
      typename std::iterator_traits<_Iterator>::difference_type;

  _CXX14_CONSTEXPR _fft_impl::coef<value_type> c;
  auto __h = __builtin_ctz(std::distance(__first, __last));

  if _CXX17_CONSTEXPR (_Inverse) {
    for (difference_type __p = 1; __p >> __h ^ 1; __p <<= 1) {
      value_type __iw = 1;
      auto __l = __first;

      for (auto __i = 1 << __h; __l != __last;
           __iw *= c.is[__builtin_ctz(--__i)]) {
        auto __r = std::next(__l, __p);

        for (auto __mid = __r; __l != __mid; ++__l, ++__r) {
          auto __tmp = (*__l - *__r) * __iw;
          *__l += *__r;
          *__r = __tmp;
        }

        __l = __r;
      }
    }

    while (__first != __last) *--__last *= c.ip2[__h];
  }

  else {
    for (difference_type __p = 1 << __h; __p >>= 1;) {
      value_type __w = -1;
      auto __l = __first;

      for (auto __i = 1 << __h; __l != __last;
           __w *= c.s[__builtin_ctz(--__i)]) {
        auto __r = std::next(__l, __p);

        for (auto __mid = __r; __l != __mid; ++__l, ++__r) {
          auto __tmp = *__l;
          *__l -= *__r *= __w;
          *__r += __tmp;
        }

        __l = __r;
      }
    }
  }
}

template <class _Iterator>
void fft(_Iterator __first, std::size_t __n) noexcept {
  fft(__first, std::next(__first, __n));
}

template <class _Iterator>
void ifft(_Iterator __first, _Iterator __last) noexcept {
  fft<true>(__first, __last);
}

template <class _Iterator>
void ifft(_Iterator __first, std::size_t __n) noexcept {
  ifft(__first, std::next(__first, __n));
}

template <size_t _Nm, size_t _Dm, class _Container, class _Index>
decltype(auto) access(_Container &__a, const _Index &__i) {
  if _CXX17_CONSTEXPR (_Nm != _Dm)
    return access<_Nm + 1, _Dm>(__a[__i[_Nm]], __i);
  else
    return __a;
}

template <bool _Inverse, size_t _Dm, class _Container, class _Tp, class _Index>
void dive(_Container &__a, const _Tp &__t, _Index &__i) {
  if _CXX17_CONSTEXPR (has_size<_Tp>::value) {
    for (__i.emplace_back(0); __i.back() != std::size(__t); ++__i.back())
      dive<_Inverse, _Dm + 1>(__a, __t[__i.back()], __i);
    __i.pop_back();
  }

  else {
    static std::vector<_Tp> __work;
    // Resize to a power of 2.

    size_t __len = 1 << (31 - __builtin_clz(std::size(__a)));
    if (__work.size() < __len) __work.resize(__len);

    for (size_t __k = 0; __k != __len; ++__k)
      __work[__k] = std::move(access<0, _Dm>(__a[__k], __i));

    fft<_Inverse>(__work.data(), __work.data() + __len);

    for (size_t __k = 0; __k != __len; ++__k)
      access<0, _Dm>(__a[__k], __i) = std::move(__work[__k]);
  }
}

template <bool _Inverse, class _Container> void fft(_Container &__a) {
  if _CXX17_CONSTEXPR (has_size<_Container>::value) {
    if _CXX17_CONSTEXPR (has_resize<_Container>::value)
      // Resize to a power of 2.
      __a.resize(1 << (32 - __builtin_clz(__a.size() - 1)));

    std::vector<size_t> __i;
    dive<_Inverse, 0>(__a, __a[0], __i);

    for (size_t __k = 0; __k != std::size(__a); ++__k) fft<_Inverse>(__a[__k]);
  }
}

template <class _Container> auto conv_resize(_Container &__a, _Container &__b) {
  std::array<size_t, get_dimension<_Container>::value> __s;
  rec(__a, __s);
  rec(__b, __s);
  return __s;
}

template <size_t _Nm, class _Container, class _Size>
void rec(const _Container &__a, _Size &__s) {
  if _CXX17_CONSTEXPR (_Nm != __s.size()) {
    __s[_Nm] = std::max(__s[_Nm], std::size(__a));
    for (auto &__x : __a) rec<_Nm + 1>(__x, __s);
  }
}

}  // namespace workspace

#line 2 "src/algebra/fft.hpp"

/**
 * @file fft.hpp
 * @brief Fast Fourier Transform
 */

#include <vector>

#line 2 "src/algebra/complex.hpp"

/**
 * @file complex.hpp
 * @brief Complex Number
 */

namespace workspace {

// Complex number.
template <class _Tp> class complex {
  _Tp re, im;

  friend constexpr complex conj(const complex &x) noexcept {
    return {x.re, -x.im};
  }

  friend constexpr _Tp abs(const complex &x) noexcept {
    return hypot(x.re, x.im);
  }

  friend constexpr _Tp arg(const complex &x) noexcept {
    return atan2(x.re, x.im);
  }

  template <class _Is>
  friend constexpr _Is &operator>>(_Is &__is, complex &x) noexcept {
    return __is >> x.re >> x.im;
  }

  template <class _Os>
  friend constexpr _Os &operator<<(_Os &__os, const complex &x) noexcept {
    return __os << x.re << ' ' << x.im;
  }

 public:
  constexpr complex() noexcept : re{}, im{} {}
  constexpr complex(_Tp _re) noexcept : re{_re}, im{} {}
  constexpr complex(_Tp _re, _Tp _im) noexcept : re{_re}, im{_im} {}

  constexpr _Tp real() const noexcept { return re; }
  constexpr void real(_Tp _re) noexcept { re = _re; }

  constexpr _Tp imag() const noexcept { return im; }
  constexpr void imag(_Tp _im) noexcept { im = _im; }

  constexpr complex operator+() const noexcept { return *this; }
  constexpr complex operator-() const noexcept { return {-re, -im}; }

  constexpr complex &operator+=(const complex &x) noexcept {
    return re += x.re, im += x.im, *this;
  }
  constexpr complex &operator-=(const complex &x) noexcept {
    return re -= x.re, im -= x.im, *this;
  }
  constexpr complex &operator*=(const complex &x) noexcept {
    _Tp _re{re * x.re - im * x.im};
    return im = im * x.re + x.im * re, re = _re, *this;
  }
  constexpr complex &operator*=(_Tp x) noexcept {
    return re *= x, im *= x, *this;
  }
  constexpr complex &operator/=(const complex &x) noexcept {
    return (*this *= conj(x)) /= re * re + im * im;
  }
  constexpr complex &operator/=(_Tp x) noexcept {
    return re /= x, im /= x, *this;
  }

  constexpr complex operator+(const complex &x) const noexcept {
    return {re + x.re, im + x.im};
  }
  constexpr complex operator-(const complex &x) const noexcept {
    return {re - x.re, im - x.im};
  }

  constexpr complex operator*(const complex &x) const noexcept {
    return complex(*this) *= x;
  }

  constexpr complex operator*(_Tp x) const noexcept { return {re * x, im * x}; }

  constexpr complex operator/(const complex &x) const noexcept {
    return complex(*this) /= x;
  }

  constexpr complex operator/(_Tp x) const noexcept { return {re / x, im / x}; }
};

}  // namespace workspace
#line 2 "lib/cxx17"

#line 2 "lib/cxx14"

#ifndef _CXX14_CONSTEXPR
#if __cplusplus >= 201402L
#define _CXX14_CONSTEXPR constexpr
#else
#define _CXX14_CONSTEXPR
#endif
#endif
#line 4 "lib/cxx17"

#ifndef _CXX17_CONSTEXPR
#if __cplusplus >= 201703L
#define _CXX17_CONSTEXPR constexpr
#else
#define _CXX17_CONSTEXPR
#endif
#endif

#ifndef _CXX17_STATIC_ASSERT
#if __cplusplus >= 201703L
#define _CXX17_STATIC_ASSERT static_assert
#else
#define _CXX17_STATIC_ASSERT assert
#endif
#endif

#include <iterator>

#if __cplusplus < 201703L

namespace std {

/**
 *  @brief  Return the size of a container.
 *  @param  __cont  Container.
 */
template <typename _Container>
constexpr auto size(const _Container& __cont) noexcept(noexcept(__cont.size()))
    -> decltype(__cont.size()) {
  return __cont.size();
}

/**
 *  @brief  Return the size of an array.
 */
template <typename _Tp, size_t _Nm>
constexpr size_t size(const _Tp (&)[_Nm]) noexcept {
  return _Nm;
}

/**
 *  @brief  Return whether a container is empty.
 *  @param  __cont  Container.
 */
template <typename _Container>
[[nodiscard]] constexpr auto empty(const _Container& __cont) noexcept(
    noexcept(__cont.empty())) -> decltype(__cont.empty()) {
  return __cont.empty();
}

/**
 *  @brief  Return whether an array is empty (always false).
 */
template <typename _Tp, size_t _Nm>
[[nodiscard]] constexpr bool empty(const _Tp (&)[_Nm]) noexcept {
  return false;
}

/**
 *  @brief  Return whether an initializer_list is empty.
 *  @param  __il  Initializer list.
 */
template <typename _Tp>
[[nodiscard]] constexpr bool empty(initializer_list<_Tp> __il) noexcept {
  return __il.size() == 0;
}

struct monostate {};

}  // namespace std

#else

#include <variant>

#endif
#line 2 "src/number_theory/ext_gcd.hpp"

/**
 * @file ext_gcd.hpp
 * @brief Extended Euclidean Algorithm
 */

#include <tuple>

#line 2 "src/utils/sfinae.hpp"

/**
 * @file sfinae.hpp
 * @brief SFINAE
 */

#include <cstdint>
#line 10 "src/utils/sfinae.hpp"
#include <type_traits>

#ifndef __INT128_DEFINED__

#ifdef __SIZEOF_INT128__
#define __INT128_DEFINED__ 1
#else
#define __INT128_DEFINED__ 0
#endif

#endif

namespace std {

#if __INT128_DEFINED__

template <> struct make_signed<__uint128_t> { using type = __int128_t; };
template <> struct make_signed<__int128_t> { using type = __int128_t; };

template <> struct make_unsigned<__uint128_t> { using type = __uint128_t; };
template <> struct make_unsigned<__int128_t> { using type = __uint128_t; };

template <> struct is_signed<__uint128_t> : std::false_type {};
template <> struct is_signed<__int128_t> : std::true_type {};

template <> struct is_unsigned<__uint128_t> : std::true_type {};
template <> struct is_unsigned<__int128_t> : std::false_type {};

#endif

}  // namespace std

namespace workspace {

template <class Tp, class... Args> struct variadic_front { using type = Tp; };

template <class... Args> struct variadic_back;

template <class Tp> struct variadic_back<Tp> { using type = Tp; };

template <class Tp, class... Args> struct variadic_back<Tp, Args...> {
  using type = typename variadic_back<Args...>::type;
};

template <class type, template <class> class trait>
using enable_if_trait_type = typename std::enable_if<trait<type>::value>::type;

/**
 * @brief Return type of subscripting ( @c [] ) access.
 */
template <class _Tp>
using subscripted_type =
    typename std::decay<decltype(std::declval<_Tp&>()[0])>::type;

template <class Container>
using element_type = typename std::decay<decltype(*std::begin(
    std::declval<Container&>()))>::type;

template <class _Tp, class = void> struct has_begin : std::false_type {};

template <class _Tp>
struct has_begin<
    _Tp, std::__void_t<decltype(std::begin(std::declval<const _Tp&>()))>>
    : std::true_type {
  using type = decltype(std::begin(std::declval<const _Tp&>()));
};

template <class _Tp, class = void> struct has_size : std::false_type {};

template <class _Tp>
struct has_size<_Tp, std::__void_t<decltype(std::size(std::declval<_Tp>()))>>
    : std::true_type {};

template <class _Tp, class = void> struct has_resize : std::false_type {};

template <class _Tp>
struct has_resize<_Tp, std::__void_t<decltype(std::declval<_Tp>().resize(
                           std::declval<size_t>()))>> : std::true_type {};

template <class _Tp, class = void> struct has_mod : std::false_type {};

template <class _Tp>
struct has_mod<_Tp, std::__void_t<decltype(_Tp::mod)>> : std::true_type {};

template <class _Tp, class = void> struct is_integral_ext : std::false_type {};
template <class _Tp>
struct is_integral_ext<
    _Tp, typename std::enable_if<std::is_integral<_Tp>::value>::type>
    : std::true_type {};

#if __INT128_DEFINED__

template <> struct is_integral_ext<__int128_t> : std::true_type {};
template <> struct is_integral_ext<__uint128_t> : std::true_type {};

#endif

#if __cplusplus >= 201402

template <class _Tp>
constexpr static bool is_integral_ext_v = is_integral_ext<_Tp>::value;

#endif

template <typename _Tp, typename = void> struct multiplicable_uint {
  using type = uint_least32_t;
};
template <typename _Tp>
struct multiplicable_uint<
    _Tp,
    typename std::enable_if<(2 < sizeof(_Tp)) &&
                            (!__INT128_DEFINED__ || sizeof(_Tp) <= 4)>::type> {
  using type = uint_least64_t;
};

#if __INT128_DEFINED__

template <typename _Tp>
struct multiplicable_uint<_Tp,
                          typename std::enable_if<(4 < sizeof(_Tp))>::type> {
  using type = __uint128_t;
};

#endif

template <typename _Tp> struct multiplicable_int {
  using type =
      typename std::make_signed<typename multiplicable_uint<_Tp>::type>::type;
};

template <typename _Tp> struct multiplicable {
  using type = std::conditional_t<
      is_integral_ext<_Tp>::value,
      std::conditional_t<std::is_signed<_Tp>::value,
                         typename multiplicable_int<_Tp>::type,
                         typename multiplicable_uint<_Tp>::type>,
      _Tp>;
};

template <class> struct first_arg { using type = void; };

template <class _R, class _Tp, class... _Args>
struct first_arg<_R(_Tp, _Args...)> {
  using type = _Tp;
};

template <class _R, class _Tp, class... _Args>
struct first_arg<_R (*)(_Tp, _Args...)> {
  using type = _Tp;
};

template <class _G, class _R, class _Tp, class... _Args>
struct first_arg<_R (_G::*)(_Tp, _Args...)> {
  using type = _Tp;
};

template <class _G, class _R, class _Tp, class... _Args>
struct first_arg<_R (_G::*)(_Tp, _Args...) const> {
  using type = _Tp;
};

template <class _Tp, class = void> struct parse_compare : first_arg<_Tp> {};

template <class _Tp>
struct parse_compare<_Tp, std::__void_t<decltype(&_Tp::operator())>>
    : first_arg<decltype(&_Tp::operator())> {};

template <class _Container, class = void> struct get_dimension {
  static constexpr size_t value = 0;
};

template <class _Container>
struct get_dimension<_Container,
                     std::enable_if_t<has_begin<_Container>::value>> {
  static constexpr size_t value =
      1 + get_dimension<typename std::iterator_traits<
              typename has_begin<_Container>::type>::value_type>::value;
};

}  // namespace workspace
#line 11 "src/number_theory/ext_gcd.hpp"

namespace workspace {

/**
 * @param __a Integer
 * @param __b Integer
 * @return Pair of integers (x, y) s.t. ax + by = g = gcd(a, b) and (b = 0 or 0
 * <= x < |b/g|) and (a = 0 or -|a/g| < y <= 0). Return (0, 0) if (a, b) = (0,
 * 0).
 */
template <typename _T1, typename _T2>
constexpr auto ext_gcd(_T1 __a, _T2 __b) noexcept {
  static_assert(is_integral_ext<_T1>::value);
  static_assert(is_integral_ext<_T2>::value);

  using value_type = typename std::make_signed<
      typename std::common_type<_T1, _T2>::type>::type;
  using result_type = std::pair<value_type, value_type>;

  value_type a{__a}, b{__b}, p{1}, q{}, r{}, s{1};

  while (b != value_type(0)) {
    auto t = a / b;
    r ^= p ^= r ^= p -= t * r;
    s ^= q ^= s ^= q -= t * s;
    b ^= a ^= b ^= a -= t * b;
  }

  if (a < 0) p = -p, q = -q, a = -a;

  if (p < 0) {
    __a /= a, __b /= a;

    if (__b > 0)
      p += __b, q -= __a;
    else
      p -= __b, q += __a;
  }

  return result_type{p, q};
}

/**
 * @param __a Integer
 * @param __b Integer
 * @param __c Integer
 * @return Pair of integers (x, y) s.t. ax + by = c and (b = 0 or 0 <= x <
 * |b/g|). Return (0, 0) if there is no solution.
 */
template <typename _T1, typename _T2, typename _T3>
constexpr auto ext_gcd(_T1 __a, _T2 __b, _T3 __c) noexcept {
  static_assert(is_integral_ext<_T1>::value);
  static_assert(is_integral_ext<_T2>::value);
  static_assert(is_integral_ext<_T3>::value);

  using value_type = typename std::make_signed<
      typename std::common_type<_T1, _T2, _T3>::type>::type;
  using result_type = std::pair<value_type, value_type>;

  value_type a{__a}, b{__b}, p{1}, q{}, r{}, s{1};

  while (b != value_type(0)) {
    auto t = a / b;
    r ^= p ^= r ^= p -= t * r;
    s ^= q ^= s ^= q -= t * s;
    b ^= a ^= b ^= a -= t * b;
  }

  if (__c % a) return result_type{};

  __a /= a, __b /= a, __c /= a;
  p *= __c, q *= __c;

  if (__b != value_type(0)) {
    auto t = p / __b;
    p -= __b * t;
    q += __a * t;

    if (p < 0) {
      if (__b > 0)
        p += __b, q -= __a;
      else
        p -= __b, q += __a;
    }
  }

  return result_type{p, q};
}

}  // namespace workspace
#line 2 "src/number_theory/primitive_root.hpp"

/**
 * @file primitive_root.hpp
 * @brief Primitive Root
 * @date 2020-12-28
 */

#line 10 "src/number_theory/primitive_root.hpp"

namespace workspace {

/**
 * @brief Compile time primitive root.
 *
 * @tparam __mod Positive integer
 * @return Minimum positive one if it exists. Otherwise 0.
 */
template <class Tp>
constexpr typename std::enable_if<(is_integral_ext<Tp>::value), Tp>::type
primitive_root(const Tp __mod) noexcept {
  assert(__mod > 0);
  using int_type = typename multiplicable_uint<Tp>::type;

  int_type __r = __mod, __p[16] = {}, *__q = __p;
  for (int_type __i = 2; __i <= __r / __i; ++__i) {
    if (__r % __i) continue;
    *__q++ = __i;
    while (!(__r % __i)) __r /= __i;
  }
  if (__r != 1) *__q++ = __r;

  int_type __tot = __mod;
  for (__q = __p; *__q; *__q++ = 0) (__tot /= *__q) *= *__q - 1;
  __r = __tot, __q = __p + 1, __p[0] = 1;
  for (int_type __i = 2; __i <= __r / __i; ++__i) {
    if (__r % __i) continue;
    *__q++ = __i;
    while (!(__r % __i)) __r /= __i;
  }
  if (__r != 1) *__q++ = __r;

  for (Tp __r = 1; __r != __mod; ++__r) {
    auto __cnt = 0;
    for (__q = __p; *__q; ++__q) {
      int_type __w = 1;
      for (int_type __e = __tot / *__q, __x = __r; __e;
           __e >>= 1, (__x *= __x) %= __mod)
        if (__e & 1) (__w *= __x) %= __mod;
      if (__w == 1 && ++__cnt > 1) break;
    }
    if (__cnt == 1) return __r;
  }

  return 0;
};

}  // namespace workspace
#line 15 "src/algebra/fft.hpp"

namespace workspace {

namespace _fft_impl {

template <class _Tp, bool = std::is_floating_point<_Tp>::value, class = void>
struct to_float {
  using type = double;
};

template <class _Tp> struct to_float<_Tp, true> { using type = _Tp; };

// template <class _Tp>
// struct to_float<_Tp, false, std::enable_if_t<sizeof(_Tp) <= sizeof(float)>> {
//   using type = float;
// };

template <class _Tp>
struct to_float<_Tp, false, std::enable_if_t<(sizeof(_Tp) > sizeof(double))>> {
  using type = long double;
};

// Assume ntt-friendly mod.
template <class _Tp> struct field {
  using type = std::conditional_t<has_mod<_Tp>::value, _Tp,
                                  complex<typename to_float<_Tp>::type>>;
};

template <class _Tp> struct field<complex<_Tp>> : field<_Tp> {};

// Modular
template <class _Tp, int _Nm = 29, bool = has_mod<_Tp>::value> struct coef {
  _Tp s[_Nm], is[_Nm], ip2[_Nm];

  _CXX14_CONSTEXPR coef() : s{}, is{}, ip2{1, (1 + _Tp::mod) / 2} {
    if (_Tp::mod < 2) return;

    int cnt2 = std::min(__builtin_ctz(_Tp::mod - 1), _Nm + 1);

    _Tp e = 1;
    _Tp w = primitive_root(_Tp::mod);
    for (auto p = (_Tp::mod - 1) >> cnt2; p; p >>= 1, w *= w)
      if (p & 1) e *= w;

    _Tp ie = ext_gcd(decltype(_Tp::mod)(e), _Tp::mod).first;
    _Tp es[_Nm]{}, ies[_Nm]{};

    for (int i = cnt2; i >= 2; i--) {
      es[i - 2] = e, e *= e;
      ies[i - 2] = ie, ie *= ie;
    }

    e = ie = 1;
    for (int i = 0; i < cnt2 - 1; i++) {
      s[i] = es[i] * e, e *= ies[i];
      is[i] = ies[i] * ie, ie *= es[i];
    }

    for (int i = 1; i < _Nm - 1; ++i) ip2[i + 1] = ip2[i] * ip2[1];
  }
};

// Complex
template <class _Tp, int _Nm> struct coef<_Tp, _Nm, false> {
  _Tp s[_Nm], is[_Nm], ip2[_Nm];

  static_assert(_Nm < 30);

  _CXX14_CONSTEXPR static _Tp es[29] = {
      {0, 1},
      {0.70710678118654752438189403651, 0.70710678118654752443610414514},
      {0.92387953251128675610142140795, 0.38268343236508977172325753068},
      {0.98078528040323044911909938781, 0.19509032201612826785692544201},
      {0.99518472667219688623102546998, 0.09801714032956060199569840382},
      {0.99879545620517239270077028412, 0.04906767432741801425693899119},
      {0.99969881869620422009748220149, 0.02454122852291228803212346128},
      {0.99992470183914454093764001552, 0.01227153828571992607945510345},
      {0.99998117528260114264494415325, 0.00613588464915447535972750246},
      {0.99999529380957617150137498041, 0.00306795676296597627029751672},
      {0.99999882345170190993313003025, 0.00153398018628476561237225788},
      {0.99999970586288221914474799723, 0.00076699031874270452695124765},
      {0.99999992646571785113833452651, 0.00038349518757139558906815188},
      {0.99999998161642929381167504976, 0.00019174759731070330743679009},
      {0.99999999540410731290905263501, 0.00009587379909597734587360460},
      {0.99999999885102682753608427379, 0.00004793689960306688454884772},
      {0.99999999971275670682981095982, 0.00002396844980841821872882467},
      {0.99999999992818917670745273995, 0.00001198422490506970642183282},
      {0.99999999998204729416331065783, 0.00000599211245264242784278378},
      {0.99999999999551182356793271877, 0.00000299605622633466075058210},
      {0.99999999999887795586487812538, 0.00000149802811316901122883643},
      {0.99999999999971948897977205850, 0.00000074901405658471572113723},
      {0.99999999999992987223139048746, 0.00000037450702829238412391495},
      {0.99999999999998246807140014902, 0.00000018725351414619534486931},
      {0.99999999999999561700429751010, 0.00000009362675707309808280024},
      {0.99999999999999890425107437752, 0.00000004681337853654909269501},
      {0.99999999999999972607632112153, 0.00000002340668926827455275977},
      {0.99999999999999993153263280754, 0.00000001170334463413727718121},
      {0.99999999999999998286960567472, 0.00000000585167231706863869077}};

  _CXX14_CONSTEXPR coef() : s{}, is{}, ip2{1, .5} {
    _Tp ies[_Nm];
    for (int i = 0; i < _Nm; ++i) ies[i] = _Tp(1) / es[i];

    _Tp e = 1, ie = 1;
    for (int i = 0; i < _Nm; i++) {
      s[i] = es[i] * e, e *= ies[i];
      is[i] = ies[i] * ie, ie *= es[i];
    }

    for (int i = 1; i < _Nm - 1; ++i) ip2[i + 1] = ip2[i] * ip2[1];
  }
};

}  // namespace _fft_impl

template <bool _Inverse = false, class _Iterator>
void fft(_Iterator __first, _Iterator __last) noexcept {
  using value_type = typename std::iterator_traits<_Iterator>::value_type;
  using difference_type =
      typename std::iterator_traits<_Iterator>::difference_type;

  _CXX14_CONSTEXPR _fft_impl::coef<value_type> c;
  auto __h = __builtin_ctz(std::distance(__first, __last));

  if _CXX17_CONSTEXPR (_Inverse) {
    for (difference_type __p = 1; __p >> __h ^ 1; __p <<= 1) {
      value_type __iw = 1;
      auto __l = __first;

      for (auto __i = 1 << __h; __l != __last;
           __iw *= c.is[__builtin_ctz(--__i)]) {
        auto __r = std::next(__l, __p);

        for (auto __mid = __r; __l != __mid; ++__l, ++__r) {
          auto __tmp = (*__l - *__r) * __iw;
          *__l += *__r;
          *__r = __tmp;
        }

        __l = __r;
      }
    }

    while (__first != __last) *--__last *= c.ip2[__h];
  }

  else {
    for (difference_type __p = 1 << __h; __p >>= 1;) {
      value_type __w = -1;
      auto __l = __first;

      for (auto __i = 1 << __h; __l != __last;
           __w *= c.s[__builtin_ctz(--__i)]) {
        auto __r = std::next(__l, __p);

        for (auto __mid = __r; __l != __mid; ++__l, ++__r) {
          auto __tmp = *__l;
          *__l -= *__r *= __w;
          *__r += __tmp;
        }

        __l = __r;
      }
    }
  }
}

template <class _Iterator>
void fft(_Iterator __first, std::size_t __n) noexcept {
  fft(__first, std::next(__first, __n));
}

template <class _Iterator>
void ifft(_Iterator __first, _Iterator __last) noexcept {
  fft<true>(__first, __last);
}

template <class _Iterator>
void ifft(_Iterator __first, std::size_t __n) noexcept {
  ifft(__first, std::next(__first, __n));
}

template <size_t _Nm, size_t _Dm, class _Container, class _Index>
decltype(auto) access(_Container &__a, const _Index &__i) {
  if _CXX17_CONSTEXPR (_Nm != _Dm)
    return access<_Nm + 1, _Dm>(__a[__i[_Nm]], __i);
  else
    return __a;
}

template <bool _Inverse, size_t _Dm, class _Container, class _Tp, class _Index>
void dive(_Container &__a, const _Tp &__t, _Index &__i) {
  if _CXX17_CONSTEXPR (has_size<_Tp>::value) {
    for (__i.emplace_back(0); __i.back() != std::size(__t); ++__i.back())
      dive<_Inverse, _Dm + 1>(__a, __t[__i.back()], __i);
    __i.pop_back();
  }

  else {
    static std::vector<_Tp> __work;
    // Resize to a power of 2.

    size_t __len = 1 << (31 - __builtin_clz(std::size(__a)));
    if (__work.size() < __len) __work.resize(__len);

    for (size_t __k = 0; __k != __len; ++__k)
      __work[__k] = std::move(access<0, _Dm>(__a[__k], __i));

    fft<_Inverse>(__work.data(), __work.data() + __len);

    for (size_t __k = 0; __k != __len; ++__k)
      access<0, _Dm>(__a[__k], __i) = std::move(__work[__k]);
  }
}

template <bool _Inverse, class _Container> void fft(_Container &__a) {
  if _CXX17_CONSTEXPR (has_size<_Container>::value) {
    if _CXX17_CONSTEXPR (has_resize<_Container>::value)
      // Resize to a power of 2.
      __a.resize(1 << (32 - __builtin_clz(__a.size() - 1)));

    std::vector<size_t> __i;
    dive<_Inverse, 0>(__a, __a[0], __i);

    for (size_t __k = 0; __k != std::size(__a); ++__k) fft<_Inverse>(__a[__k]);
  }
}

template <class _Container> auto conv_resize(_Container &__a, _Container &__b) {
  std::array<size_t, get_dimension<_Container>::value> __s;
  rec(__a, __s);
  rec(__b, __s);
  return __s;
}

template <size_t _Nm, class _Container, class _Size>
void rec(const _Container &__a, _Size &__s) {
  if _CXX17_CONSTEXPR (_Nm != __s.size()) {
    __s[_Nm] = std::max(__s[_Nm], std::size(__a));
    for (auto &__x : __a) rec<_Nm + 1>(__x, __s);
  }
}

}  // namespace workspace

Fast Fourier Transform (src/algebra/fft.hpp)

Depends on

Required by

Verified with

Code

Fast Fourier Transform
(src/algebra/fft.hpp)