kth root(Tonelli-Shanks algorithm)
(modulo/mod-kth-root.hpp)

• View this file on GitHub
• Last update: 2024-03-04 16:48:10+09:00
• Include: #include "modulo/mod-kth-root.hpp"

kth root

整数$a,k$、素数$p$に対して$x^k \equiv a \mod p$を満たす$x$を$\mathrm{O}(\min (p,k)^{\frac{1}{4}})$で計算するライブラリ。

概要

37zigenさんの解説を参考にしました。

原始根を利用して考察する。$\mod p$上の原始根を$r$と置き、$x\equiv r^y,a\equiv r^b$とすると上の等式は

\[x^k \equiv a \mod p \leftrightarrow ky \equiv b \mod p-1\]

と言い換えられる。ここで$g = \gcd(k,p-1)$と置くと、上の合同式が成り立つ条件は$b \equiv 0 \mod g$であることがわかる。よって、

• $b \not \equiv 0 \mod g$の場合は解が存在しない。
• $b \equiv 0 \mod g$の場合は、$\frac{k}{g}$と$\frac{p-1}{g}$が互いに素なので$\left(\frac{k}{g}\right)^{-1}\mod \frac{p-1}{g}$が存在して、この時$y \equiv \frac{b}{g} \left(\frac{k}{g}\right)^{-1}\mod \frac{p-1}{g}$が解となる。

以上より解を求めることが出来たが、このままでは計算量は離散対数問題と同じ$\mathrm{O}(\sqrt{p})$なので、mod sqrtの時と同様にTonelli-Shanks algorithmを利用して高速化を図る。

まず、解の存在条件は

\[b \equiv 0 \mod g\] \[\leftrightarrow \frac{b(p-1)}{g}\equiv 0 \mod p-1\] \[\leftrightarrow a^\frac{p-1}{g} \equiv 1 \mod p\]

となる。以下は解が存在する時のみ考える。

$k\left(gk^{-1} \mod \frac{p-1}{g}\right)\equiv g\mod p-1$であるから、与式は

\[x^k \equiv a \mod p \leftrightarrow x^g\equiv \mathrm{pow}\left(a,gk^{-1} \mod \frac{p-1}{g}\right)\mod p\]

となる。(このように$x$の底を$k$から$g$に変換することで、$g\mid p-1$の性質を利用したToneli-Shanksのアルゴリズムが適用可能になる。)ここで$g = \Pi_i p_i^{e_i}$と置いたとき、

\[z^{p_i^{e_i}}\equiv c \mod p \ \ \ \cdots (\ast)\]

を解くことが出来れば、($\ast$)を繰り返し元の式に適用することで$x$を求めることが出来る。よって($\ast$)を高速に計算する方法を考えればよく、この問題はmod sqrtと同様にしてTonelli-Shanksのアルゴリズムが適用出来る。

原始根$r$に対して$r^w=z,r^d=c,p-1=sp_i^t$($\gcd(s,p_i)=1$)とおくと、

\[(\ast) \leftrightarrow p_i^{e_i}w\equiv d\mod p\] \[\leftrightarrow p_i^{e_i}w\equiv d\mod s\ \wedge p_i^{e_i} w\equiv d\mod p_i^t\]

となる。ここで、$\left(-s^{-1} \mod {p_i}^{e_i}\right)=u$と置き、$z$の初期値を$z_0 \equiv \mathrm{pow}(c,\frac{su+1}{ {p_i}^{e_i}}) \mod p$とおくと、

\[p_i^{e_i}w_0 \equiv d(su+1) \mod p-1\] \[\leftrightarrow p_i^{e_i}w_0\equiv d\mod s\ \wedge p_i^{e_i} w\equiv d(su+1)\mod p_i^t\]

となるので$\mod s$の方は条件を満たしており、$\mod p_i^{t_i}$を調整すればよいとわかる。

ここで、$\mathrm{pow}\left(v,\frac{p-1}{p_i}\right)\not\equiv 1$である$v$を乱択で選ぶ。(条件を満たす$v$は$1-\frac{1}{p_i}$の確率で手に入る。)この時、

\[\mathrm{Ind}_r v^{s} \equiv 0\mod s\ ,\mathrm{Ind}_r v^s \not\equiv 0\mod p_i\]

が成り立つ。($\mathrm{Ind}_r$は$r$を底、$p$を法としたときの指数。)この$v^s$を利用してTonelli-Shanks algorithmを実行する。

まず、現在の$z$の誤差項$p_i^{e^i} z - d\mod p_i^{e_i}$の誤差が$np_i^e$と表せるとする。この時$e$は、$c^{-1}z^{p_i^{e_i}}$を$p_i$乗していって$1$になるまでにかかった回数を$t’$としたときに $e=t-t’$となる。$e$が分かったら、$z^{ {p_i}^{t’-1}}\equiv 1$を満たすまで$z$に$\mathrm{pow}(v^s,p^{e-e_i})$を掛け続ける。このアルゴリズムを$t’=0$になるまで繰り返せば$z$を求めることが出来る。

このままだとTonelli-Shanksのループ一回あたり最大$p_i-1$回の乗算が必要となるが、最後の$\mathrm{pow}(v^s,2^{e-e_i})$を掛ける所でBaby Step Giant Stepを利用することで、ループあたりの乗算回数を$\mathrm{O}(\sqrt{p_i})$回に落とすことが出来る。

また、Tonelli-Shanksのループ回数は高々$\lfloor\log_{p_i}p\rfloor - 1$回となる。つまり、$t=e_i=1$の時($g$と$p-1$がともに$p_i$で1回ずつしか割り切れないとき)はループに入らないので、$p_i \gt \sqrt{p-1}$である素数に対しては高速に解が計算できるとわかる。よって、全体の計算量は$g=\gcd(k,p-1)$が$\mathrm{O}(\sqrt{p})$程度の重複度2の素因数を含む時が最大で、この時$\mathrm{O}(g^{\frac{1}{4}})=\mathrm{O}(\min (p,k)^{\frac{1}{4}})$となる。

$g$の素因数分解もMillar-rabinとPollardの$\rho$法を利用すれば$\mathrm{O}(g^{\frac{1}{4}})$で計算できるので、全体の計算量は$\mathrm{O}(\min (p,k)^{\frac{1}{4}})$となる。

Depends on

• internal/internal-math.hpp
• internal/internal-seed.hpp
• internal/internal-type-traits.hpp
• misc/rng.hpp
• modint/arbitrary-montgomery-modint.hpp
• 高速素因数分解(Miller Rabin/Pollard's Rho) (prime/fast-factorize.hpp)
• Miller-Rabin primality test (prime/miller-rabin.hpp)

Verified with

• verify/verify-yosupo-math/yosupo-kth-root-mod.test.cpp

Code

#pragma once

#include "../internal/internal-math.hpp"
#include "../modint/arbitrary-montgomery-modint.hpp"
#include "../prime/fast-factorize.hpp"

namespace kth_root_mod {

// fast BS-GS
template <typename T>
struct Memo {
  Memo(const T &g, int s, int period)
      : size(1 << __lg(min(s, period))),
        mask(size - 1),
        period(period),
        vs(size),
        os(size + 1) {
    T x(1);
    for (int i = 0; i < size; ++i, x *= g) os[x.get() & mask]++;
    for (int i = 1; i < size; ++i) os[i] += os[i - 1];
    x = 1;
    for (int i = 0; i < size; ++i, x *= g) vs[--os[x.get() & mask]] = {x, i};
    gpow = x;
    os[size] = size;
  }
  int find(T x) const {
    for (int t = 0; t < period; t += size, x *= gpow) {
      for (int m = (x.get() & mask), i = os[m]; i < os[m + 1]; ++i) {
        if (x == vs[i].first) {
          int ret = vs[i].second - t;
          return ret < 0 ? ret + period : ret;
        }
      }
    }
    assert(0);
  }
  T gpow;
  int size, mask, period;
  vector<pair<T, int>> vs;
  vector<int> os;
};

template <typename INT, typename LINT, typename mint>
mint pe_root(INT c, INT pi, INT ei, INT p) {
  if (mint::get_mod() != decltype(mint::a)(p)) mint::set_mod(p);
  INT s = p - 1, t = 0;
  while (s % pi == 0) s /= pi, ++t;
  INT pe = 1;
  for (INT _ = 0; _ < ei; ++_) pe *= pi;

  INT u = internal::inv(pe - s % pe, pe);
  mint mc = c, one = 1;
  mint z = mc.pow((s * u + 1) / pe);
  mint zpe = mc.pow(s * u);
  if (zpe == one) return z;
  assert(t > ei);

  mint vs;
  {
    INT ptm1 = 1;
    for (INT _ = 0; _ < t - 1; ++_) ptm1 *= pi;
    for (mint v = 2;; v += one) {
      vs = v.pow(s);
      if (vs.pow(ptm1) != one) break;
    }
  }

  mint vspe = vs.pow(pe);
  INT vs_e = ei;
  mint base = vspe;
  for (INT _ = 0; _ < t - ei - 1; _++) base = base.pow(pi);
  Memo<mint> memo(base, (INT)(sqrt(t - ei) * sqrt(pi)) + 1, pi);

  while (zpe != one) {
    mint tmp = zpe;
    INT td = 0;
    while (tmp != 1) ++td, tmp = tmp.pow(pi);
    INT e = t - td;
    while (vs_e != e) {
      vs = vs.pow(pi);
      vspe = vspe.pow(pi);
      ++vs_e;
    }

    // BS-GS ... find (zpe * ( vspe ^ n ) ) ^( p_i ^ (td - 1) ) = 1
    mint base_zpe = zpe.inverse();
    for (INT _ = 0; _ < td - 1; _++) base_zpe = base_zpe.pow(pi);
    INT bsgs = memo.find(base_zpe);

    z *= vs.pow(bsgs);
    zpe *= vspe.pow(bsgs);
  }
  return z;
}

template <typename INT, typename LINT, typename mint>
INT inner_kth_root(INT a, INT k, INT p) {
  a %= p;
  if (k == 0) return a == 1 ? a : -1;
  if (a <= 1 || k <= 1) return a;
  assert(p > 2);
  if (mint::get_mod() != decltype(mint::a)(p)) mint::set_mod(p);
  INT g = gcd(p - 1, k);
  if (internal::modpow<INT, LINT>(a, (p - 1) / g, p) != 1) return -1;
  a = mint(a).pow(internal::inv(k / g, (p - 1) / g)).get();
  unordered_map<INT, int> fac;
  for (auto &f : factorize(g)) fac[f]++;
  if (mint::get_mod() != decltype(mint::a)(p)) mint::set_mod(p);
  for (auto pp : fac)
    a = pe_root<INT, LINT, mint>(a, pp.first, pp.second, p).get();
  return a;
}

int64_t kth_root(int64_t a, int64_t k, int64_t p) {
  if (max({a, k, p}) < (1LL << 30))
    return inner_kth_root<int32_t, int64_t,
                          ArbitraryLazyMontgomeryModInt<163553130>>(a, k, p);
  else
    return inner_kth_root<int64_t, __int128_t,
                          ArbitraryLazyMontgomeryModInt64bit<504025646>>(a, k,
                                                                         p);
}

}  // namespace kth_root_mod
using kth_root_mod::kth_root;

/**
 * @brief kth root(Tonelli-Shanks algorithm)
 * @docs docs/modulo/mod-kth-root.md
 */

#line 2 "modulo/mod-kth-root.hpp"

#line 2 "internal/internal-math.hpp"

#line 2 "internal/internal-type-traits.hpp"

#include <type_traits>
using namespace std;

namespace internal {
template <typename T>
using is_broadly_integral =
    typename conditional_t<is_integral_v<T> || is_same_v<T, __int128_t> ||
                               is_same_v<T, __uint128_t>,
                           true_type, false_type>::type;

template <typename T>
using is_broadly_signed =
    typename conditional_t<is_signed_v<T> || is_same_v<T, __int128_t>,
                           true_type, false_type>::type;

template <typename T>
using is_broadly_unsigned =
    typename conditional_t<is_unsigned_v<T> || is_same_v<T, __uint128_t>,
                           true_type, false_type>::type;

#define ENABLE_VALUE(x) \
  template <typename T> \
  constexpr bool x##_v = x<T>::value;

ENABLE_VALUE(is_broadly_integral);
ENABLE_VALUE(is_broadly_signed);
ENABLE_VALUE(is_broadly_unsigned);
#undef ENABLE_VALUE

#define ENABLE_HAS_TYPE(var)                                   \
  template <class, class = void>                               \
  struct has_##var : false_type {};                            \
  template <class T>                                           \
  struct has_##var<T, void_t<typename T::var>> : true_type {}; \
  template <class T>                                           \
  constexpr auto has_##var##_v = has_##var<T>::value;

#define ENABLE_HAS_VAR(var)                                     \
  template <class, class = void>                                \
  struct has_##var : false_type {};                             \
  template <class T>                                            \
  struct has_##var<T, void_t<decltype(T::var)>> : true_type {}; \
  template <class T>                                            \
  constexpr auto has_##var##_v = has_##var<T>::value;

}  // namespace internal
#line 4 "internal/internal-math.hpp"

namespace internal {

#include <cassert>
#include <utility>
#include <vector>
using namespace std;

// a mod p
template <typename T>
T safe_mod(T a, T p) {
  a %= p;
  if constexpr (is_broadly_signed_v<T>) {
    if (a < 0) a += p;
  }
  return a;
}

// 返り値：pair(g, x)
// s.t. g = gcd(a, b), xa = g (mod b), 0 <= x < b/g
template <typename T>
pair<T, T> inv_gcd(T a, T p) {
  static_assert(is_broadly_signed_v<T>);
  a = safe_mod(a, p);
  if (a == 0) return {p, 0};
  T b = p, x = 1, y = 0;
  while (a != 0) {
    T q = b / a;
    swap(a, b %= a);
    swap(x, y -= q * x);
  }
  if (y < 0) y += p / b;
  return {b, y};
}

// 返り値 : a^{-1} mod p
// gcd(a, p) != 1 が必要
template <typename T>
T inv(T a, T p) {
  static_assert(is_broadly_signed_v<T>);
  a = safe_mod(a, p);
  T b = p, x = 1, y = 0;
  while (a != 0) {
    T q = b / a;
    swap(a, b %= a);
    swap(x, y -= q * x);
  }
  assert(b == 1);
  return y < 0 ? y + p : y;
}

// T : 底の型
// U : T*T がオーバーフローしない かつ 指数の型
template <typename T, typename U>
T modpow(T a, U n, T p) {
  a = safe_mod(a, p);
  T ret = 1 % p;
  while (n != 0) {
    if (n % 2 == 1) ret = U(ret) * a % p;
    a = U(a) * a % p;
    n /= 2;
  }
  return ret;
}

// 返り値 : pair(rem, mod)
// 解なしのときは {0, 0} を返す
template <typename T>
pair<T, T> crt(const vector<T>& r, const vector<T>& m) {
  static_assert(is_broadly_signed_v<T>);
  assert(r.size() == m.size());
  int n = int(r.size());
  T r0 = 0, m0 = 1;
  for (int i = 0; i < n; i++) {
    assert(1 <= m[i]);
    T r1 = safe_mod(r[i], m[i]), m1 = m[i];
    if (m0 < m1) swap(r0, r1), swap(m0, m1);
    if (m0 % m1 == 0) {
      if (r0 % m1 != r1) return {0, 0};
      continue;
    }
    auto [g, im] = inv_gcd(m0, m1);
    T u1 = m1 / g;
    if ((r1 - r0) % g) return {0, 0};
    T x = (r1 - r0) / g % u1 * im % u1;
    r0 += x * m0;
    m0 *= u1;
    if (r0 < 0) r0 += m0;
  }
  return {r0, m0};
}

}  // namespace internal
#line 2 "modint/arbitrary-montgomery-modint.hpp"

#include <iostream>
using namespace std;

template <typename Int, typename UInt, typename Long, typename ULong, int id>
struct ArbitraryLazyMontgomeryModIntBase {
  using mint = ArbitraryLazyMontgomeryModIntBase;

  inline static UInt mod;
  inline static UInt r;
  inline static UInt n2;
  static constexpr int bit_length = sizeof(UInt) * 8;

  static UInt get_r() {
    UInt ret = mod;
    while (mod * ret != 1) ret *= UInt(2) - mod * ret;
    return ret;
  }
  static void set_mod(UInt m) {
    assert(m < (UInt(1u) << (bit_length - 2)));
    assert((m & 1) == 1);
    mod = m, n2 = -ULong(m) % m, r = get_r();
  }
  UInt a;

  ArbitraryLazyMontgomeryModIntBase() : a(0) {}
  ArbitraryLazyMontgomeryModIntBase(const Long &b)
      : a(reduce(ULong(b % mod + mod) * n2)){};

  static UInt reduce(const ULong &b) {
    return (b + ULong(UInt(b) * UInt(-r)) * mod) >> bit_length;
  }

  mint &operator+=(const mint &b) {
    if (Int(a += b.a - 2 * mod) < 0) a += 2 * mod;
    return *this;
  }
  mint &operator-=(const mint &b) {
    if (Int(a -= b.a) < 0) a += 2 * mod;
    return *this;
  }
  mint &operator*=(const mint &b) {
    a = reduce(ULong(a) * b.a);
    return *this;
  }
  mint &operator/=(const mint &b) {
    *this *= b.inverse();
    return *this;
  }

  mint operator+(const mint &b) const { return mint(*this) += b; }
  mint operator-(const mint &b) const { return mint(*this) -= b; }
  mint operator*(const mint &b) const { return mint(*this) *= b; }
  mint operator/(const mint &b) const { return mint(*this) /= b; }

  bool operator==(const mint &b) const {
    return (a >= mod ? a - mod : a) == (b.a >= mod ? b.a - mod : b.a);
  }
  bool operator!=(const mint &b) const {
    return (a >= mod ? a - mod : a) != (b.a >= mod ? b.a - mod : b.a);
  }
  mint operator-() const { return mint(0) - mint(*this); }
  mint operator+() const { return mint(*this); }

  mint pow(ULong n) const {
    mint ret(1), mul(*this);
    while (n > 0) {
      if (n & 1) ret *= mul;
      mul *= mul, n >>= 1;
    }
    return ret;
  }

  friend ostream &operator<<(ostream &os, const mint &b) {
    return os << b.get();
  }

  friend istream &operator>>(istream &is, mint &b) {
    Long t;
    is >> t;
    b = ArbitraryLazyMontgomeryModIntBase(t);
    return (is);
  }

  mint inverse() const {
    Int x = get(), y = get_mod(), u = 1, v = 0;
    while (y > 0) {
      Int t = x / y;
      swap(x -= t * y, y);
      swap(u -= t * v, v);
    }
    return mint{u};
  }

  UInt get() const {
    UInt ret = reduce(a);
    return ret >= mod ? ret - mod : ret;
  }

  static UInt get_mod() { return mod; }
};

// id に適当な乱数を割り当てて使う
template <int id>
using ArbitraryLazyMontgomeryModInt =
    ArbitraryLazyMontgomeryModIntBase<int, unsigned int, long long,
                                      unsigned long long, id>;
template <int id>
using ArbitraryLazyMontgomeryModInt64bit =
    ArbitraryLazyMontgomeryModIntBase<long long, unsigned long long, __int128_t,
                                      __uint128_t, id>;
#line 2 "prime/fast-factorize.hpp"

#include <cstdint>
#include <numeric>
#line 6 "prime/fast-factorize.hpp"
using namespace std;

#line 2 "misc/rng.hpp"

#line 2 "internal/internal-seed.hpp"

#include <chrono>
using namespace std;

namespace internal {
unsigned long long non_deterministic_seed() {
  unsigned long long m =
      chrono::duration_cast<chrono::nanoseconds>(
          chrono::high_resolution_clock::now().time_since_epoch())
          .count();
  m ^= 9845834732710364265uLL;
  m ^= m << 24, m ^= m >> 31, m ^= m << 35;
  return m;
}
unsigned long long deterministic_seed() { return 88172645463325252UL; }

// 64 bit の seed 値を生成 (手元では seed 固定)
// 連続で呼び出すと同じ値が何度も返ってくるので注意
// #define RANDOMIZED_SEED するとシードがランダムになる
unsigned long long seed() {
#if defined(NyaanLocal) && !defined(RANDOMIZED_SEED)
  return deterministic_seed();
#else
  return non_deterministic_seed();
#endif
}

}  // namespace internal
#line 4 "misc/rng.hpp"

namespace my_rand {
using i64 = long long;
using u64 = unsigned long long;

// [0, 2^64 - 1)
u64 rng() {
  static u64 _x = internal::seed();
  return _x ^= _x << 7, _x ^= _x >> 9;
}

// [l, r]
i64 rng(i64 l, i64 r) {
  assert(l <= r);
  return l + rng() % u64(r - l + 1);
}

// [l, r)
i64 randint(i64 l, i64 r) {
  assert(l < r);
  return l + rng() % u64(r - l);
}

// choose n numbers from [l, r) without overlapping
vector<i64> randset(i64 l, i64 r, i64 n) {
  assert(l <= r && n <= r - l);
  unordered_set<i64> s;
  for (i64 i = n; i; --i) {
    i64 m = randint(l, r + 1 - i);
    if (s.find(m) != s.end()) m = r - i;
    s.insert(m);
  }
  vector<i64> ret;
  for (auto& x : s) ret.push_back(x);
  return ret;
}

// [0.0, 1.0)
double rnd() { return rng() * 5.42101086242752217004e-20; }
// [l, r)
double rnd(double l, double r) {
  assert(l < r);
  return l + rnd() * (r - l);
}

template <typename T>
void randshf(vector<T>& v) {
  int n = v.size();
  for (int i = 1; i < n; i++) swap(v[i], v[randint(0, i + 1)]);
}

}  // namespace my_rand

using my_rand::randint;
using my_rand::randset;
using my_rand::randshf;
using my_rand::rnd;
using my_rand::rng;
#line 2 "prime/miller-rabin.hpp"

#line 4 "prime/miller-rabin.hpp"
using namespace std;

#line 8 "prime/miller-rabin.hpp"

namespace fast_factorize {

template <typename T, typename U>
bool miller_rabin(const T& n, vector<T> ws) {
  if (n <= 2) return n == 2;
  if (n % 2 == 0) return false;

  T d = n - 1;
  while (d % 2 == 0) d /= 2;
  U e = 1, rev = n - 1;
  for (T w : ws) {
    if (w % n == 0) continue;
    T t = d;
    U y = internal::modpow<T, U>(w, t, n);
    while (t != n - 1 && y != e && y != rev) y = y * y % n, t *= 2;
    if (y != rev && t % 2 == 0) return false;
  }
  return true;
}

bool miller_rabin_u64(unsigned long long n) {
  return miller_rabin<unsigned long long, __uint128_t>(
      n, {2, 325, 9375, 28178, 450775, 9780504, 1795265022});
}

template <typename mint>
bool miller_rabin(unsigned long long n, vector<unsigned long long> ws) {
  if (n <= 2) return n == 2;
  if (n % 2 == 0) return false;

  if (mint::get_mod() != n) mint::set_mod(n);
  unsigned long long d = n - 1;
  while (~d & 1) d >>= 1;
  mint e = 1, rev = n - 1;
  for (unsigned long long w : ws) {
    if (w % n == 0) continue;
    unsigned long long t = d;
    mint y = mint(w).pow(t);
    while (t != n - 1 && y != e && y != rev) y *= y, t *= 2;
    if (y != rev && t % 2 == 0) return false;
  }
  return true;
}

bool is_prime(unsigned long long n) {
  using mint32 = ArbitraryLazyMontgomeryModInt<96229631>;
  using mint64 = ArbitraryLazyMontgomeryModInt64bit<622196072>;

  if (n <= 2) return n == 2;
  if (n % 2 == 0) return false;
  if (n < (1uLL << 30)) {
    return miller_rabin<mint32>(n, {2, 7, 61});
  } else if (n < (1uLL << 62)) {
    return miller_rabin<mint64>(
        n, {2, 325, 9375, 28178, 450775, 9780504, 1795265022});
  } else {
    return miller_rabin_u64(n);
  }
}

}  // namespace fast_factorize

using fast_factorize::is_prime;

/**
 * @brief Miller-Rabin primality test
 */
#line 12 "prime/fast-factorize.hpp"

namespace fast_factorize {
using u64 = uint64_t;

template <typename mint, typename T>
T pollard_rho(T n) {
  if (~n & 1) return 2;
  if (is_prime(n)) return n;
  if (mint::get_mod() != n) mint::set_mod(n);
  mint R, one = 1;
  auto f = [&](mint x) { return x * x + R; };
  auto rnd_ = [&]() { return rng() % (n - 2) + 2; };
  while (1) {
    mint x, y, ys, q = one;
    R = rnd_(), y = rnd_();
    T g = 1;
    constexpr int m = 128;
    for (int r = 1; g == 1; r <<= 1) {
      x = y;
      for (int i = 0; i < r; ++i) y = f(y);
      for (int k = 0; g == 1 && k < r; k += m) {
        ys = y;
        for (int i = 0; i < m && i < r - k; ++i) q *= x - (y = f(y));
        g = gcd(q.get(), n);
      }
    }
    if (g == n) do
        g = gcd((x - (ys = f(ys))).get(), n);
      while (g == 1);
    if (g != n) return g;
  }
  exit(1);
}

using i64 = long long;

vector<i64> inner_factorize(u64 n) {
  using mint32 = ArbitraryLazyMontgomeryModInt<452288976>;
  using mint64 = ArbitraryLazyMontgomeryModInt64bit<401243123>;

  if (n <= 1) return {};
  u64 p;
  if (n <= (1LL << 30)) {
    p = pollard_rho<mint32, uint32_t>(n);
  } else if (n <= (1LL << 62)) {
    p = pollard_rho<mint64, uint64_t>(n);
  } else {
    exit(1);
  }
  if (p == n) return {i64(p)};
  auto l = inner_factorize(p);
  auto r = inner_factorize(n / p);
  copy(begin(r), end(r), back_inserter(l));
  return l;
}

vector<i64> factorize(u64 n) {
  auto ret = inner_factorize(n);
  sort(begin(ret), end(ret));
  return ret;
}

map<i64, i64> factor_count(u64 n) {
  map<i64, i64> mp;
  for (auto &x : factorize(n)) mp[x]++;
  return mp;
}

vector<i64> divisors(u64 n) {
  if (n == 0) return {};
  vector<pair<i64, i64>> v;
  for (auto &p : factorize(n)) {
    if (v.empty() || v.back().first != p) {
      v.emplace_back(p, 1);
    } else {
      v.back().second++;
    }
  }
  vector<i64> ret;
  auto f = [&](auto rc, int i, i64 x) -> void {
    if (i == (int)v.size()) {
      ret.push_back(x);
      return;
    }
    rc(rc, i + 1, x);
    for (int j = 0; j < v[i].second; j++) rc(rc, i + 1, x *= v[i].first);
  };
  f(f, 0, 1);
  sort(begin(ret), end(ret));
  return ret;
}

}  // namespace fast_factorize

using fast_factorize::divisors;
using fast_factorize::factor_count;
using fast_factorize::factorize;

/**
 * @brief 高速素因数分解(Miller Rabin/Pollard's Rho)
 * @docs docs/prime/fast-factorize.md
 */
#line 6 "modulo/mod-kth-root.hpp"

namespace kth_root_mod {

// fast BS-GS
template <typename T>
struct Memo {
  Memo(const T &g, int s, int period)
      : size(1 << __lg(min(s, period))),
        mask(size - 1),
        period(period),
        vs(size),
        os(size + 1) {
    T x(1);
    for (int i = 0; i < size; ++i, x *= g) os[x.get() & mask]++;
    for (int i = 1; i < size; ++i) os[i] += os[i - 1];
    x = 1;
    for (int i = 0; i < size; ++i, x *= g) vs[--os[x.get() & mask]] = {x, i};
    gpow = x;
    os[size] = size;
  }
  int find(T x) const {
    for (int t = 0; t < period; t += size, x *= gpow) {
      for (int m = (x.get() & mask), i = os[m]; i < os[m + 1]; ++i) {
        if (x == vs[i].first) {
          int ret = vs[i].second - t;
          return ret < 0 ? ret + period : ret;
        }
      }
    }
    assert(0);
  }
  T gpow;
  int size, mask, period;
  vector<pair<T, int>> vs;
  vector<int> os;
};

template <typename INT, typename LINT, typename mint>
mint pe_root(INT c, INT pi, INT ei, INT p) {
  if (mint::get_mod() != decltype(mint::a)(p)) mint::set_mod(p);
  INT s = p - 1, t = 0;
  while (s % pi == 0) s /= pi, ++t;
  INT pe = 1;
  for (INT _ = 0; _ < ei; ++_) pe *= pi;

  INT u = internal::inv(pe - s % pe, pe);
  mint mc = c, one = 1;
  mint z = mc.pow((s * u + 1) / pe);
  mint zpe = mc.pow(s * u);
  if (zpe == one) return z;
  assert(t > ei);

  mint vs;
  {
    INT ptm1 = 1;
    for (INT _ = 0; _ < t - 1; ++_) ptm1 *= pi;
    for (mint v = 2;; v += one) {
      vs = v.pow(s);
      if (vs.pow(ptm1) != one) break;
    }
  }

  mint vspe = vs.pow(pe);
  INT vs_e = ei;
  mint base = vspe;
  for (INT _ = 0; _ < t - ei - 1; _++) base = base.pow(pi);
  Memo<mint> memo(base, (INT)(sqrt(t - ei) * sqrt(pi)) + 1, pi);

  while (zpe != one) {
    mint tmp = zpe;
    INT td = 0;
    while (tmp != 1) ++td, tmp = tmp.pow(pi);
    INT e = t - td;
    while (vs_e != e) {
      vs = vs.pow(pi);
      vspe = vspe.pow(pi);
      ++vs_e;
    }

    // BS-GS ... find (zpe * ( vspe ^ n ) ) ^( p_i ^ (td - 1) ) = 1
    mint base_zpe = zpe.inverse();
    for (INT _ = 0; _ < td - 1; _++) base_zpe = base_zpe.pow(pi);
    INT bsgs = memo.find(base_zpe);

    z *= vs.pow(bsgs);
    zpe *= vspe.pow(bsgs);
  }
  return z;
}

template <typename INT, typename LINT, typename mint>
INT inner_kth_root(INT a, INT k, INT p) {
  a %= p;
  if (k == 0) return a == 1 ? a : -1;
  if (a <= 1 || k <= 1) return a;
  assert(p > 2);
  if (mint::get_mod() != decltype(mint::a)(p)) mint::set_mod(p);
  INT g = gcd(p - 1, k);
  if (internal::modpow<INT, LINT>(a, (p - 1) / g, p) != 1) return -1;
  a = mint(a).pow(internal::inv(k / g, (p - 1) / g)).get();
  unordered_map<INT, int> fac;
  for (auto &f : factorize(g)) fac[f]++;
  if (mint::get_mod() != decltype(mint::a)(p)) mint::set_mod(p);
  for (auto pp : fac)
    a = pe_root<INT, LINT, mint>(a, pp.first, pp.second, p).get();
  return a;
}

int64_t kth_root(int64_t a, int64_t k, int64_t p) {
  if (max({a, k, p}) < (1LL << 30))
    return inner_kth_root<int32_t, int64_t,
                          ArbitraryLazyMontgomeryModInt<163553130>>(a, k, p);
  else
    return inner_kth_root<int64_t, __int128_t,
                          ArbitraryLazyMontgomeryModInt64bit<504025646>>(a, k,
                                                                         p);
}

}  // namespace kth_root_mod
using kth_root_mod::kth_root;

/**
 * @brief kth root(Tonelli-Shanks algorithm)
 * @docs docs/modulo/mod-kth-root.md
 */

Back to top page