例 · 一份真正“现代”的 FFT 模板

4月 05, 2021

欣闻 UOJ 已经用上了新 GCC，安排上了 C++20，我便拿出之前玩耍 MSVC 时写的 FFT 板子，改了改一些地方，交交试试。

结果上比较顺利，我觉得至少里面用上的 Concepts 和各类编译期求值计数还是有一定参考价值，并且这份 FFT 的优点还是比较明显的：

包含一个 ModInt 类，作用如其名，模板参数是数字类型和模数。支持模意义下需要的各类运算符（位运算和大小比较自然无意义就没写了）和流输入输出；
包含一个 FFT 类，用于做 DFT 和 IDFT，模板参数是数组元素类型，目前支持两种：ModInt<A,B> 和 std::complex<double>，分别对应 NTT 和 FFT。不需要指定，自动根据模板参数选择算法；
NTT 不用自己找模数，在编译期能自动求出最小的原根来使用，找不到原根会报错；
所有的非法模板参数都会在编译期较友好地提示，NTT模数对应的长度被超出时也会在运行期出错，理所当然地，用了 1e9+7 这类非 NTT 模数也是可以找出错误告诉你的；
本来在 MSVC 上还用了 Modules 实现辅助计算类和变换类分离，现在提供 GCC 版本自然是删掉了。

缺点也是有一些的，直接列出来比较好

最大的缺点，内部采用 std::vector 实现，速度较慢；
采用工程和算法混合的阴间风格命名，部分地方比较混乱，如果你看的不爽可以自己改改。

具体来说，你可以在里面学习：

Concepts 的基本用法，require 套套套；
简单 SFINAE 和 Type Traits；
Constexpr 求值，进阶：std::array 等；
If Constexpr 和 Static Assertion 并用，让 #define 从地球上消失的技巧；
合理使用引用和重载减少大型数据复制。

这么些人人都会的技术。

没事啦，权当玩玩x

UOJ 34 提交记录：

NTT: 1342ms/7612kb/C++ 20/7.7kb
FFT: 458ms/15160kb/C++ 20/7.7kb

FFT 好快啊，毕竟原数组值域只有 9 呢x

通用模板代码如下：

#include <iostream>
#include <vector>
#include <array>
#include <type_traits>
#include <complex>
#include <concepts>

namespace zmile {

template<class T>
concept GroupElement =
std::is_arithmetic_v<T>
|| requires(T x) { x = 1; x* x; };

template<class T>
concept BinaryPowerNumber = requires(T x) {
  x >>= 1;
  x & 1;
};

template<GroupElement ValueT, BinaryPowerNumber PowerT>
constexpr ValueT qpow(ValueT a, PowerT b, ValueT unit = 1) {
  for (; b; b >>= 1, a = a * a)
    if (b & 1) unit = unit * a;
  return unit;
}

template <class ValueT = int, ValueT M = 998244353>
requires(std::is_integral_v<ValueT>&& M > 0 && requires(ValueT x) { x = M; })
class ModInt {
  using ThisT = ModInt<ValueT, M>;
  ValueT value;
public:
  using value_t = ValueT;

  constexpr ModInt() : value(0) {  }
  constexpr ModInt(ValueT value) : value((value % M + M) % M) {  }
  //constexpr ModInt(int64_t value) : value((value % M + M) % M) {  }

  inline constexpr explicit operator bool() const { return value; }

  inline constexpr ThisT operator-(void) const { return (M - value) % M; }

  inline constexpr ThisT operator+(const ThisT& rhs) const { return (value + rhs.value) % M; }
  inline constexpr ThisT operator-(const ThisT& rhs) const { return *this + -rhs; }
  inline constexpr ThisT operator*(const ThisT& rhs) const { return ValueT(1ull * value * rhs.value % M); }
  inline constexpr ThisT operator/(const ThisT& rhs) const { return *this * qpow(rhs, M-2); }

  inline constexpr ThisT& operator+=(const ThisT& rhs) { return *this = *this + rhs; }
  inline constexpr ThisT& operator-=(const ThisT& rhs) { return *this = *this - rhs; }
  inline constexpr ThisT& operator*=(const ThisT& rhs) { return *this = *this * rhs; }
  inline constexpr ThisT& operator/=(const ThisT& rhs) { return *this = *this / rhs; }

  inline constexpr bool operator==(const ThisT& rhs) const { return value == rhs.get_value(); }
  inline constexpr bool operator==(const ValueT& rhs) const { return value == rhs; }

  inline constexpr ThisT get_inv() const { return ThisT(1) / value; }
  inline constexpr ValueT get_value() const { return value; }
};

template <class ValueT, ValueT M>
requires(std::is_integral_v<ValueT>&& M > 0 && requires(ValueT x) { x = M; })
std::ostream& operator<<(std::ostream& os, ModInt<ValueT, M> val)
{
  return os << val.get_value();
}

template <class ValueT, ValueT M>
requires(std::is_integral_v<ValueT>&& M > 0 && requires(ValueT x) { x = M; })
std::istream& operator>>(std::istream& is, ModInt<ValueT, M>& val)
{
  static ValueT temp_val;
  is >> temp_val;
  val = temp_val;
  return is;
}

// requires mod is prime
template<auto mod>
requires(std::is_integral_v<decltype(mod)>
  && requires() { mod % mod; mod - mod; mod / mod; qpow(mod, mod); })
constexpr decltype(mod) find_primitive_root_internal() {
  using ValueT = decltype(mod);
  // this is enough for int64_t
  constexpr int MAX_PRIME_FACTOR_COUNT = 17;
  std::array<ValueT, MAX_PRIME_FACTOR_COUNT> fac{ 0 };
  //std::vector<ValueT> fac;
  auto vm = mod - 1; int pos = 0;
  for (int x = 2; x * x <= vm; x++) {
    if (vm % x == 0) {
      fac[pos++] = x;
      while (vm % x == 0) vm /= x;
    }
  }
  if (vm != 1) fac[pos++] = vm;

  for (auto ans = 2; ans < mod; ans++) {
    bool failed = false;
    for (auto x : fac) {
      if (x == 0 || x == 1) break;
      if (qpow<class ModInt<ValueT, mod>>(ans, (mod - 1) / x) == 1) {
        failed = true;
        break;
      }
    }
    if (!failed) return ans;
  }

  return 0;
}

template <auto mod>
constexpr auto find_primitive_root() {
  constexpr auto rt = find_primitive_root_internal<mod>();
  static_assert(rt, "FATAL ERR: the primitive root doesn't exists!");
  return rt;
}

// NTT max valid poly length
inline constexpr static auto get_max_valid_length(auto M) {
  auto Mt = M; size_t cnt2 = 0;
  while (!(Mt & 1))
    Mt >>= 1, cnt2++;
  return 1ull << cnt2;
}

}


namespace zmile {

template <class T>
struct is_modint : public std::false_type {  };
template <class T, T M>
struct is_modint<ModInt<T, M>> : public std::true_type {  };
template <class T>
constexpr bool is_modint_v = is_modint<T>::value;

template <class T>
struct get_modint_mod;
template <class T, T M>
struct get_modint_mod<ModInt<T, M>> {
  static const T value = M;
};
template <class T>
constexpr auto get_modint_mod_v = get_modint_mod<T>::value;

template<class ValueT = ModInt<>>
requires(
  requires(ValueT x) {
    x + x;
    x * x;
    x / ValueT(10);
    ValueT(1) / x;
    x = 1;
  }
)
class FFT {
  size_t n;
  std::vector<size_t> rev;
  using poly_t = std::vector<ValueT>;
  const double PI = std::acos(-1);

public:
  FFT(size_t n = 0) {
    if (n) set_size(n);
  }

  enum class TransformDirection : int8_t {
    DFT = 1,
    IDFT = -1
  };

  // set an enough length to STORE THE MULT RESULT
  void set_size(size_t size) {
    for (size_t L = 1; (n = L) <= size; L <<= 1);
    rev.assign(n, 0);
    for (size_t i = 1; i < n; i++)
      rev[i] = rev[i >> 1] >> 1 | (n >> 1) * (i & 1);
  }

  inline poly_t transform(const poly_t& poly, TransformDirection dir) {
    poly_t result;
    transform(poly, result, dir);
    return result;
  }

  // also work correctly when poly === result
  inline void transform(const poly_t& poly, poly_t& result, TransformDirection dir) {
    if (poly.empty()) return result.clear();

    if (n < poly.size())
      set_size(poly.size());

    if (&result != &poly)
      result = poly;
    result.resize(n);
    if constexpr (is_modint_v<ValueT>) {
      if (get_max_valid_length(get_modint_mod_v<ValueT>-1) < n)
        throw std::logic_error("FATAL ERR: the mod prime is not strong enough to hold the NTT result!");
    }

    for (size_t i = 0; i < n; i++) {
      if (i < rev[i])
        std::swap(result[i], result[rev[i]]);
    }

    for (size_t h = 2; h <= n; h <<= 1) {

      ValueT wn, w;
      if constexpr (is_modint_v<ValueT>) {
        wn = qpow(ValueT(find_primitive_root<get_modint_mod_v<ValueT>>()),
                  (get_modint_mod_v<ValueT> - 1) / h);
      } else if constexpr (std::is_same_v<ValueT, std::complex<double>>) {
        wn = { cos(2*PI/h), sin(2*PI/h) };
      } else {
        []<bool flag = false>(){
          static_assert(flag, "FATAL ERR: FFT only support std::complex<double> and ModInt.");
        }();
      }

      if (dir == TransformDirection::IDFT)
        wn = ValueT(1) / wn;

      for (size_t j = 0; j < n; j += h) {
        w = 1;
        for (size_t k = j; k < j + (h >> 1); k++, w *= wn) {
          ValueT x = result[k], y = result[k + (h >> 1)] * w;
          result[k] = x + y, result[k + (h >> 1)] = x - y;
        }
      }
    }

    if (dir == TransformDirection::IDFT) {
      for (size_t i = 0; i < n; i++)
        result[i] /= n;
    }
  }
};

}

/// ======================
/// Use the template above
/// ======================

const int MOD = 1004535809; // alternative: 1e9+7 leads to runtime error
using num_t = std::complex<double>;// alternative: zmile::ModInt<int, MOD>;
using dir_t = zmile::FFT<num_t>::TransformDirection;

int main() {
    size_t n, m;
    num_t val;
    std::cin >> n >> m;
    
    zmile::FFT<num_t> transformer{ n + m + 1 };
    std::vector<num_t> poly1, poly2;

    for (int i = 0; i <= n; i++) {
        std::cin >> val;
        poly1.push_back(val);
    }
    for (int i = 0; i <= m; i++) {
        std::cin >> val;
        poly2.push_back(val);
    }

    transformer.transform(poly1, poly1, dir_t::DFT);
    transformer.transform(poly2, poly2, dir_t::DFT);

    for (int i = 0; i < poly1.size(); i++) {
        poly1[i] *= poly2[i];
    }

    transformer.transform(poly1, poly1, dir_t::IDFT);
    for (int i = 0; i <= n + m; i++) {
        std::cout << int(poly1[i].real() + 0.5) << ' '; // ModInt: directly output poly1[i]
    }

    return 0;
}

评论和共享

例 · 一份真正“现代”的 FFT 模板

7mile