Library

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub jellc/Library

:warning: Number Theoretic Transform
(src/algebra/ntt.hpp)

Depends on

Code

#pragma once

/**
 * @file ntt.hpp
 * @brief Number Theoretic Transform
 * @date 2021-02-20
 *
 *
 */

#include "src/number_theory/ext_gcd.hpp"
#include "src/number_theory/primitive_root.hpp"

namespace workspace {

namespace ntt_impl {

/**
 * @see
 * https://github.com/atcoder/ac-library/blob/master/atcoder/convolution.hpp
 */

template <class _Tp> struct __coef {
  _Tp sum_e[30];  // sum_e[i] = ies[0] * ... * ies[i - 1] * es[i]

  constexpr __coef() : sum_e{} {
    if (_Tp::mod < 2) return;

    int cnt2 = __builtin_ctz(_Tp::mod - 1);
    _Tp e = 1;
    {
      auto p = (_Tp::mod - 1) >> cnt2;
      _Tp w = primitive_root(_Tp::mod);
      while (p) {
        if (p & 1) e *= w;
        p >>= 1;
        w *= w;
      }
    }

    _Tp ie = ext_gcd(decltype(_Tp::mod)(e), _Tp::mod).first;
    _Tp es[30] = {}, ies[30] = {};  // es[i]^(2^(2+i)) == 1
    for (int i = cnt2; i >= 2; i--) {
      // e^(2^i) == 1
      es[i - 2] = e;
      ies[i - 2] = ie;
      e *= e;
      ie *= ie;
    }

    _Tp now = 1;
    for (int i = 0; i <= cnt2 - 2; i++) {
      sum_e[i] = es[i] * now;
      now *= ies[i];
    }
  }
};

template <class _Tp> struct __icoef {
  _Tp sum_ie[30];  // sum_e[i] = ies[0] * ... * ies[i - 1] * es[i]

  constexpr __icoef() : sum_ie{} {
    if (_Tp::mod < 2) return;

    int cnt2 = __builtin_ctz(_Tp::mod - 1);
    _Tp e = 1;
    {
      auto p = (_Tp::mod - 1) >> cnt2;
      _Tp w = primitive_root(_Tp::mod);
      while (p) {
        if (p & 1) e *= w;
        p >>= 1;
        w *= w;
      }
    }

    _Tp ie = ext_gcd(decltype(_Tp::mod)(e), _Tp::mod).first;
    _Tp es[30] = {}, ies[30] = {};  // es[i]^(2^(2+i)) == 1
    for (int i = cnt2; i >= 2; i--) {
      // e^(2^i) == 1
      es[i - 2] = e;
      ies[i - 2] = ie;
      e *= e;
      ie *= ie;
    }

    _Tp now = 1;
    for (int i = 0; i <= cnt2 - 2; i++) {
      sum_ie[i] = ies[i] * now;
      now *= es[i];
    }
  }
};

template <class _Tp> struct __ipow2 {
  _Tp __ip2[30];
  constexpr __ipow2() : __ip2{1, (1 + _Tp::mod) / 2} {
    for (size_t __i = 1; __i + 1 != std::size(__ip2); ++__i)
      __ip2[__i + 1] = __ip2[__i] * __ip2[1];
  }
};

template <class _FIter>
constexpr void ntt(_FIter __first, _FIter __last) noexcept {
  using value_type = typename std::decay<decltype(*__first)>::type;
  constexpr __coef<value_type> _;

  auto __h = __builtin_ctz(std::distance(__first, __last));

  for (ptrdiff_t __p = 1 << __h; __p >>= 1;) {
    value_type now = -1;
    auto __l = __first;

    for (size_t __s = 1 << __h; __l != __last;
         now *= _.sum_e[__builtin_ctz(--__s)]) {
      auto __r = __l + __p;

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

      __l = __r;
    }
  }
}

template <class _A> constexpr void ntt(_A &a) noexcept {
  ntt(std::begin(a), std::end(a));
}

template <class _FIter>
constexpr void intt(_FIter __first, _FIter __last) noexcept {
  using value_type = typename std::decay<decltype(*__first)>::type;
  constexpr __icoef<value_type> _;

  auto __h = __builtin_ctz(std::distance(__first, __last));

  for (ptrdiff_t __p = 1; __p >> __h ^ 1; __p <<= 1) {
    value_type inow = 1;
    auto __l = __first;

    for (size_t __s = 1 << __h; __l != __last;
         inow *= _.sum_ie[__builtin_ctz(--__s)]) {
      auto __r = __l + __p;

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

      __l = __r;
    }
  }

  constexpr __ipow2<value_type> __;
  while (__first != __last) *--__last *= __.__ip2[__h];
}  // namespace ntt_impl

template <class _A> constexpr void intt(_A &a) noexcept {
  intt(std::begin(a), std::end(a));
}

}  // namespace ntt_impl

using ntt_impl::intt;
using ntt_impl::ntt;

}  // namespace workspace
#line 2 "src/algebra/ntt.hpp"

/**
 * @file ntt.hpp
 * @brief Number Theoretic Transform
 * @date 2021-02-20
 *
 *
 */

#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>
#include <iterator>
#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 13 "src/algebra/ntt.hpp"

namespace workspace {

namespace ntt_impl {

/**
 * @see
 * https://github.com/atcoder/ac-library/blob/master/atcoder/convolution.hpp
 */

template <class _Tp> struct __coef {
  _Tp sum_e[30];  // sum_e[i] = ies[0] * ... * ies[i - 1] * es[i]

  constexpr __coef() : sum_e{} {
    if (_Tp::mod < 2) return;

    int cnt2 = __builtin_ctz(_Tp::mod - 1);
    _Tp e = 1;
    {
      auto p = (_Tp::mod - 1) >> cnt2;
      _Tp w = primitive_root(_Tp::mod);
      while (p) {
        if (p & 1) e *= w;
        p >>= 1;
        w *= w;
      }
    }

    _Tp ie = ext_gcd(decltype(_Tp::mod)(e), _Tp::mod).first;
    _Tp es[30] = {}, ies[30] = {};  // es[i]^(2^(2+i)) == 1
    for (int i = cnt2; i >= 2; i--) {
      // e^(2^i) == 1
      es[i - 2] = e;
      ies[i - 2] = ie;
      e *= e;
      ie *= ie;
    }

    _Tp now = 1;
    for (int i = 0; i <= cnt2 - 2; i++) {
      sum_e[i] = es[i] * now;
      now *= ies[i];
    }
  }
};

template <class _Tp> struct __icoef {
  _Tp sum_ie[30];  // sum_e[i] = ies[0] * ... * ies[i - 1] * es[i]

  constexpr __icoef() : sum_ie{} {
    if (_Tp::mod < 2) return;

    int cnt2 = __builtin_ctz(_Tp::mod - 1);
    _Tp e = 1;
    {
      auto p = (_Tp::mod - 1) >> cnt2;
      _Tp w = primitive_root(_Tp::mod);
      while (p) {
        if (p & 1) e *= w;
        p >>= 1;
        w *= w;
      }
    }

    _Tp ie = ext_gcd(decltype(_Tp::mod)(e), _Tp::mod).first;
    _Tp es[30] = {}, ies[30] = {};  // es[i]^(2^(2+i)) == 1
    for (int i = cnt2; i >= 2; i--) {
      // e^(2^i) == 1
      es[i - 2] = e;
      ies[i - 2] = ie;
      e *= e;
      ie *= ie;
    }

    _Tp now = 1;
    for (int i = 0; i <= cnt2 - 2; i++) {
      sum_ie[i] = ies[i] * now;
      now *= es[i];
    }
  }
};

template <class _Tp> struct __ipow2 {
  _Tp __ip2[30];
  constexpr __ipow2() : __ip2{1, (1 + _Tp::mod) / 2} {
    for (size_t __i = 1; __i + 1 != std::size(__ip2); ++__i)
      __ip2[__i + 1] = __ip2[__i] * __ip2[1];
  }
};

template <class _FIter>
constexpr void ntt(_FIter __first, _FIter __last) noexcept {
  using value_type = typename std::decay<decltype(*__first)>::type;
  constexpr __coef<value_type> _;

  auto __h = __builtin_ctz(std::distance(__first, __last));

  for (ptrdiff_t __p = 1 << __h; __p >>= 1;) {
    value_type now = -1;
    auto __l = __first;

    for (size_t __s = 1 << __h; __l != __last;
         now *= _.sum_e[__builtin_ctz(--__s)]) {
      auto __r = __l + __p;

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

      __l = __r;
    }
  }
}

template <class _A> constexpr void ntt(_A &a) noexcept {
  ntt(std::begin(a), std::end(a));
}

template <class _FIter>
constexpr void intt(_FIter __first, _FIter __last) noexcept {
  using value_type = typename std::decay<decltype(*__first)>::type;
  constexpr __icoef<value_type> _;

  auto __h = __builtin_ctz(std::distance(__first, __last));

  for (ptrdiff_t __p = 1; __p >> __h ^ 1; __p <<= 1) {
    value_type inow = 1;
    auto __l = __first;

    for (size_t __s = 1 << __h; __l != __last;
         inow *= _.sum_ie[__builtin_ctz(--__s)]) {
      auto __r = __l + __p;

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

      __l = __r;
    }
  }

  constexpr __ipow2<value_type> __;
  while (__first != __last) *--__last *= __.__ip2[__h];
}  // namespace ntt_impl

template <class _A> constexpr void intt(_A &a) noexcept {
  intt(std::begin(a), std::end(a));
}

}  // namespace ntt_impl

using ntt_impl::intt;
using ntt_impl::ntt;

}  // namespace workspace
Back to top page