mod_kth_root.hpp

SECTIONMath INCLUDEnoya/mod_kth_root.hpp

在素数模下求一个满足 \(x^k=a\) 的根，或判断不存在。

\[ \displaystyle x^k\equiv a\pmod p \]

Complexity: Time: O(sqrt(g) + sum sqrt(q e) log p), g=gcd(k,p-1), over prime factors q^e of g. Space: O(max sqrt(q e)).

AC 记录：kth_root_mod。

跳到代码 · GitHub ↗

Implementation

ReferenceSourceExpanded

当前头文件，省略 include guard；依赖见 #include。

/// @complexity Time: O(sqrt(g) + sum sqrt(q e) log p), g=gcd(k,p-1), over prime factors q^e of g.
/// Space: O(max sqrt(q e)).

#include "noya/discrete_log.hpp"

#include <cassert>
#include <cmath>
#include <cstdint>
#include <map>
#include <numeric>
#include <optional>

namespace noya {

namespace mod_kth_root_internal {

inline std::uint64_t prime_power(std::uint64_t pri, int exp) {
  std::uint64_t res = 1;
  while (exp-- > 0) {
    res *= pri;
  }
  return res;
}

/// Extract a p^exp-th root. The caller guarantees existence.
inline std::uint64_t extract_prime_power_root(std::uint64_t val,
                                              std::uint64_t p, int exp,
                                              std::uint64_t mod) {
  using factorize_internal::multiply_mod;
  using factorize_internal::power_mod;

  std::uint64_t cp = mod - 1;
  int vp = 0;
  while (cp % p == 0) {
    cp /= p;
    vp++;
  }

  const std::uint64_t pw = prime_power(p, exp);
  const std::uint64_t icp = discrete_log_internal::inverse_mod(cp % pw, pw);
  const std::uint64_t ie = std::uint64_t(
      ((static_cast<unsigned __int128>(pw - 1) * icp % pw) * cp + 1) / pw);
  std::uint64_t ans = power_mod(val, ie, mod);

  std::uint64_t nqr = 2;
  while (power_mod(nqr, (mod - 1) / p, mod) == 1) {
    nqr++;
  }
  nqr = power_mod(nqr, cp, mod);

  const int gs = int(std::sqrt(static_cast<long double>(vp - exp) * p)) + 1;
  const std::uint64_t tp = prime_power(p, vp - 1);
  const std::uint64_t ts =
      power_mod(nqr, factorize_internal::multiply_mod(gs, tp, mod - 1), mod);
  std::map<std::uint64_t, int> tbl;
  std::uint64_t cur = 1;
  for (int idx = 0; idx <= gs; idx++) {
    tbl[cur] = idx;
    cur = multiply_mod(cur, ts, mod);
  }

  const std::uint64_t is =
      discrete_log_internal::inverse_mod(power_mod(nqr, tp, mod), mod);
  for (int lev = exp; lev < vp; lev++) {
    std::uint64_t err = multiply_mod(
        discrete_log_internal::inverse_mod(power_mod(ans, pw, mod), mod), val,
        mod);
    std::uint64_t tar = power_mod(err, prime_power(p, vp - 1 - lev), mod);
    bool fou = false;
    for (int ste = 0; ste <= gs; ste++) {
      auto it = tbl.find(tar);
      if (it != tbl.end()) {
        const std::uint64_t dig =
            std::uint64_t(ste) + std::uint64_t(gs) * it->second;
        const std::uint64_t ce = factorize_internal::multiply_mod(
            dig, prime_power(p, lev - exp), mod - 1);
        ans = multiply_mod(ans, power_mod(nqr, ce, mod), mod);
        fou = true;
        break;
      }
      tar = multiply_mod(tar, is, mod);
    }
    assert(fou);
  }
  return ans;
}

} // namespace mod_kth_root_internal

/// @brief Return one x satisfying x^k = val modulo a prime, or nullopt;
/// avoids a full discrete logarithm by extracting the prime-power factors of
/// gcd(k, p - 1) independently.
inline std::optional<std::uint64_t>
mod_kth_root(std::uint64_t val, std::uint64_t exp, std::uint64_t pri) {
  assert(pri >= 2 && is_prime(pri));
  val %= pri;
  if (exp == 0) {
    return val == 1 ? std::optional<std::uint64_t>(1) : std::nullopt;
  }
  if (val == 0) {
    return 0;
  }
  if (pri == 2) {
    return val;
  }
  const std::uint64_t ord = pri - 1;
  exp %= ord;
  if (exp == 0) {
    return val == 1 ? std::optional<std::uint64_t>(1) : std::nullopt;
  }
  std::uint64_t div = std::gcd(exp, ord);
  if (factorize_internal::power_mod(val, ord / div, pri) != 1) {
    return std::nullopt;
  }

  const std::uint64_t ro = ord / div;
  if (ro > 1) {
    const std::uint64_t inv =
        discrete_log_internal::inverse_mod(exp / div % ro, ro);
    val = factorize_internal::power_mod(val, inv, pri);
  }

  for (std::uint64_t fct = 2; fct * fct <= div; fct++) {
    int cnt = 0;
    while (div % fct == 0) {
      div /= fct;
      cnt++;
    }
    if (cnt > 0) {
      val = mod_kth_root_internal::extract_prime_power_root(val, fct, cnt, pri);
    }
  }
  if (div > 1) {
    val = mod_kth_root_internal::extract_prime_power_root(val, div, 1, pri);
  }
  return val;
}

/// @brief Return the number of roots to x^k = val modulo a prime.
inline std::uint64_t mod_kth_root_count(std::uint64_t val, std::uint64_t exp,
                                        std::uint64_t pri) {
  assert(pri >= 2 && is_prime(pri));
  val %= pri;
  if (exp == 0) {
    return val == 1 ? pri : 0;
  }
  if (val == 0) {
    return 1;
  }
  auto rt = mod_kth_root(val, exp, pri);
  return rt ? std::gcd(exp, pri - 1) : 0;
}

} // namespace noya

#ifndef NOYA_MOD_KTH_ROOT_HPP
#define NOYA_MOD_KTH_ROOT_HPP 1

/// @complexity Time: O(sqrt(g) + sum sqrt(q e) log p), g=gcd(k,p-1), over prime factors q^e of g.
/// Space: O(max sqrt(q e)).

#include "noya/discrete_log.hpp"

#include <cassert>
#include <cmath>
#include <cstdint>
#include <map>
#include <numeric>
#include <optional>

namespace noya {

namespace mod_kth_root_internal {

inline std::uint64_t prime_power(std::uint64_t pri, int exp) {
  std::uint64_t res = 1;
  while (exp-- > 0) {
    res *= pri;
  }
  return res;
}

/// Extract a p^exp-th root. The caller guarantees existence.
inline std::uint64_t extract_prime_power_root(std::uint64_t val,
                                              std::uint64_t p, int exp,
                                              std::uint64_t mod) {
  using factorize_internal::multiply_mod;
  using factorize_internal::power_mod;

  std::uint64_t cp = mod - 1;
  int vp = 0;
  while (cp % p == 0) {
    cp /= p;
    vp++;
  }

  const std::uint64_t pw = prime_power(p, exp);
  const std::uint64_t icp = discrete_log_internal::inverse_mod(cp % pw, pw);
  const std::uint64_t ie = std::uint64_t(
      ((static_cast<unsigned __int128>(pw - 1) * icp % pw) * cp + 1) / pw);
  std::uint64_t ans = power_mod(val, ie, mod);

  std::uint64_t nqr = 2;
  while (power_mod(nqr, (mod - 1) / p, mod) == 1) {
    nqr++;
  }
  nqr = power_mod(nqr, cp, mod);

  const int gs = int(std::sqrt(static_cast<long double>(vp - exp) * p)) + 1;
  const std::uint64_t tp = prime_power(p, vp - 1);
  const std::uint64_t ts =
      power_mod(nqr, factorize_internal::multiply_mod(gs, tp, mod - 1), mod);
  std::map<std::uint64_t, int> tbl;
  std::uint64_t cur = 1;
  for (int idx = 0; idx <= gs; idx++) {
    tbl[cur] = idx;
    cur = multiply_mod(cur, ts, mod);
  }

  const std::uint64_t is =
      discrete_log_internal::inverse_mod(power_mod(nqr, tp, mod), mod);
  for (int lev = exp; lev < vp; lev++) {
    std::uint64_t err = multiply_mod(
        discrete_log_internal::inverse_mod(power_mod(ans, pw, mod), mod), val,
        mod);
    std::uint64_t tar = power_mod(err, prime_power(p, vp - 1 - lev), mod);
    bool fou = false;
    for (int ste = 0; ste <= gs; ste++) {
      auto it = tbl.find(tar);
      if (it != tbl.end()) {
        const std::uint64_t dig =
            std::uint64_t(ste) + std::uint64_t(gs) * it->second;
        const std::uint64_t ce = factorize_internal::multiply_mod(
            dig, prime_power(p, lev - exp), mod - 1);
        ans = multiply_mod(ans, power_mod(nqr, ce, mod), mod);
        fou = true;
        break;
      }
      tar = multiply_mod(tar, is, mod);
    }
    assert(fou);
  }
  return ans;
}

} // namespace mod_kth_root_internal

/// @brief Return one x satisfying x^k = val modulo a prime, or nullopt;
/// avoids a full discrete logarithm by extracting the prime-power factors of
/// gcd(k, p - 1) independently.
inline std::optional<std::uint64_t>
mod_kth_root(std::uint64_t val, std::uint64_t exp, std::uint64_t pri) {
  assert(pri >= 2 && is_prime(pri));
  val %= pri;
  if (exp == 0) {
    return val == 1 ? std::optional<std::uint64_t>(1) : std::nullopt;
  }
  if (val == 0) {
    return 0;
  }
  if (pri == 2) {
    return val;
  }
  const std::uint64_t ord = pri - 1;
  exp %= ord;
  if (exp == 0) {
    return val == 1 ? std::optional<std::uint64_t>(1) : std::nullopt;
  }
  std::uint64_t div = std::gcd(exp, ord);
  if (factorize_internal::power_mod(val, ord / div, pri) != 1) {
    return std::nullopt;
  }

  const std::uint64_t ro = ord / div;
  if (ro > 1) {
    const std::uint64_t inv =
        discrete_log_internal::inverse_mod(exp / div % ro, ro);
    val = factorize_internal::power_mod(val, inv, pri);
  }

  for (std::uint64_t fct = 2; fct * fct <= div; fct++) {
    int cnt = 0;
    while (div % fct == 0) {
      div /= fct;
      cnt++;
    }
    if (cnt > 0) {
      val = mod_kth_root_internal::extract_prime_power_root(val, fct, cnt, pri);
    }
  }
  if (div > 1) {
    val = mod_kth_root_internal::extract_prime_power_root(val, div, 1, pri);
  }
  return val;
}

/// @brief Return the number of roots to x^k = val modulo a prime.
inline std::uint64_t mod_kth_root_count(std::uint64_t val, std::uint64_t exp,
                                        std::uint64_t pri) {
  assert(pri >= 2 && is_prime(pri));
  val %= pri;
  if (exp == 0) {
    return val == 1 ? pri : 0;
  }
  if (val == 0) {
    return 1;
  }
  auto rt = mod_kth_root(val, exp, pri);
  return rt ? std::gcd(exp, pri - 1) : 0;
}

} // namespace noya

#endif // NOYA_MOD_KTH_ROOT_HPP

#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <map>
#include <numeric>
#include <optional>
#include <unordered_map>
#include <utility>
#include <vector>

/// @complexity Time: O(sqrt(g) + sum sqrt(q e) log p), g=gcd(k,p-1), over prime factors q^e of g.
/// Space: O(max sqrt(q e)).

/// @complexity Time: O(sqrt(m)) expected time.
/// Space: O(sqrt(m)).

/// @complexity Time: O(log^3 n) primality testing; Pollard-rho factorization is expected about O(n^(1/4)).
/// Space: O(log n) recursion and factors.

namespace noya {
namespace factorize_internal {

using u64 = std::uint64_t;
using u128 = unsigned __int128;

inline u64 multiply_mod(u64 a, u64 b, u64 mod) {
  return u64(u128(a) * b % mod);
}

inline u64 power_mod(u64 a, u64 exp, u64 mod) {
  u64 res = 1;
  while (exp > 0) {
    if (exp & 1) {
      res = multiply_mod(res, a, mod);
    }
    a = multiply_mod(a, a, mod);
    exp >>= 1;
  }
  return res;
}

inline bool miller_rabin(u64 n) {
  if (n < 2) {
    return false;
  }
  for (u64 p :
       std::array<u64, 12>{2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37}) {
    if (n % p == 0) {
      return n == p;
    }
  }
  int shf = __builtin_ctzll(n - 1);
  u64 odd = (n - 1) >> shf;
  for (u64 bas :
       std::array<u64, 7>{2, 325, 9375, 28178, 450775, 9780504, 1795265022}) {
    if (bas % n == 0) {
      continue;
    }
    u64 val = power_mod(bas % n, odd, n);
    if (val == 1 || val == n - 1) {
      continue;
    }
    bool cmp = true;
    for (int i = 1; i < shf; i++) {
      val = multiply_mod(val, val, n);
      if (val == n - 1) {
        cmp = false;
        break;
      }
    }
    if (cmp) {
      return false;
    }
  }
  return true;
}

inline u64 splitmix64(u64 &st) {
  u64 z = (st += 0x9e3779b97f4a7c15ULL);
  z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
  z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
  return z ^ (z >> 31);
}

inline u64 pollard_rho(u64 n) {
  if (n % 2 == 0) {
    return 2;
  }
  if (n % 3 == 0) {
    return 3;
  }
  static u64 st = 0x123456789abcdef0ULL;
  while (true) {
    u64 y = splitmix64(st) % (n - 1) + 1;
    u64 c = splitmix64(st) % (n - 1) + 1;
    constexpr u64 blk = 128;
    u64 g = 1;
    u64 r = 1;
    u64 q = 1;
    u64 x = 0;
    u64 sy = 0;
    auto nxt = [&](u64 val) {
      return u64((u128(multiply_mod(val, val, n)) + c) % n);
    };
    while (g == 1) {
      x = y;
      for (u64 i = 0; i < r; i++) {
        y = nxt(y);
      }
      for (u64 off = 0; off < r && g == 1; off += blk) {
        sy = y;
        for (u64 i = 0; i < std::min(blk, r - off); i++) {
          y = nxt(y);
          u64 dif = x > y ? x - y : y - x;
          q = multiply_mod(q, dif, n);
        }
        g = std::gcd(q, n);
      }
      r <<= 1;
    }
    if (g == n) {
      do {
        sy = nxt(sy);
        u64 dif = x > sy ? x - sy : sy - x;
        g = std::gcd(dif, n);
      } while (g == 1);
    }
    if (g != n) {
      return g;
    }
  }
}

inline void collect_factors(u64 n, std::vector<u64> &res) {
  if (n == 1) {
    return;
  }
  if (miller_rabin(n)) {
    res.push_back(n);
    return;
  }
  u64 fct = pollard_rho(n);
  collect_factors(fct, res);
  collect_factors(n / fct, res);
}

} // namespace factorize_internal

/// @brief Deterministic Miller-Rabin primality test for unsigned 64-bit
/// integers.
inline bool is_prime(std::uint64_t n) {
  return factorize_internal::miller_rabin(n);
}

/// @brief Return the prime factors of n with multiplicity in increasing order.
inline std::vector<std::uint64_t> prime_factors(std::uint64_t n) {
  assert(n >= 1);
  std::vector<std::uint64_t> res;
  factorize_internal::collect_factors(n, res);
  std::sort(res.begin(), res.end());
  return res;
}

/// @brief Return the prime factorization of n as (prime, exponent) pairs.
inline std::vector<std::pair<std::uint64_t, int>> factorize(std::uint64_t n) {
  std::vector<std::pair<std::uint64_t, int>> res;
  for (std::uint64_t p : prime_factors(n)) {
    if (res.empty() || res.back().first != p) {
      res.emplace_back(p, 1);
    } else {
      res.back().second++;
    }
  }
  return res;
}

} // namespace noya

namespace noya {
namespace discrete_log_internal {

using u64 = std::uint64_t;
using i128 = __int128;

inline u64 inverse_mod(u64 val, u64 mod) {
  assert(mod >= 2 && std::gcd(val, mod) == 1);
  i128 c0 = 1;
  i128 cf = 0;
  u64 r0 = val;
  u64 rem = mod;
  while (rem != 0) {
    u64 quo = r0 / rem;
    u64 r1 = r0 - quo * rem;
    r0 = rem;
    rem = r1;
    i128 c1 = c0 - i128(quo) * cf;
    c0 = cf;
    cf = c1;
  }
  c0 %= i128(mod);
  if (c0 < 0) {
    c0 += mod;
  }
  return u64(c0);
}

inline u64 ceil_sqrt(u64 val) {
  u64 res = u64(std::sqrt(static_cast<long double>(val)));
  while (res != 0 && res > val / res) {
    res--;
  }
  while (res == 0 || res <= val / res) {
    if (res * res == val) {
      return res;
    }
    res++;
  }
  return res;
}

inline std::optional<u64> coprime_bsgs(u64 bas, u64 tar, u64 mod) {
  using factorize_internal::multiply_mod;
  using factorize_internal::power_mod;
  assert(mod >= 2 && std::gcd(bas, mod) == 1);
  u64 blk = ceil_sqrt(mod);
  std::unordered_map<u64, u64> bs;
  bs.reserve(std::size_t(blk * 2 + 1));
  u64 val = 1;
  for (u64 exp = 0; exp < blk; exp++) {
    bs.emplace(val, exp);
    val = multiply_mod(val, bas, mod);
  }
  u64 ib = power_mod(inverse_mod(bas, mod), blk, mod);
  val = tar;
  for (u64 gia = 0; gia <= blk; gia++) {
    auto it = bs.find(val);
    if (it != bs.end()) {
      return gia * blk + it->second;
    }
    val = multiply_mod(val, ib, mod);
  }
  return std::nullopt;
}

} // namespace discrete_log_internal

/// @brief Find the smallest x >= 0 satisfying bas^x = tar (mod mod),
/// or nullopt if no solution exists.
inline std::optional<std::uint64_t>
discrete_log(std::uint64_t bas, std::uint64_t tar, std::uint64_t mod) {
  using discrete_log_internal::coprime_bsgs;
  using factorize_internal::multiply_mod;
  assert(mod >= 1);
  if (mod == 1) {
    return 0;
  }
  bas %= mod;
  tar %= mod;
  if (tar == 1) {
    return 0;
  }

  std::uint64_t rmv = 0;
  std::uint64_t acc = 1;
  while (true) {
    std::uint64_t div = std::gcd(bas, mod);
    if (div == 1) {
      break;
    }
    if (tar == acc) {
      return rmv;
    }
    if (tar % div != 0) {
      return std::nullopt;
    }
    tar /= div;
    mod /= div;
    acc = multiply_mod(acc, bas / div, mod);
    rmv++;
  }
  if (tar == acc) {
    return rmv;
  }
  std::uint64_t nrm =
      multiply_mod(tar, discrete_log_internal::inverse_mod(acc, mod), mod);
  auto res = coprime_bsgs(bas % mod, nrm, mod);
  if (!res) {
    return std::nullopt;
  }
  return *res + rmv;
}

/// @brief Extended BSGS: find the smallest x >= 0 satisfying
/// bas^x = tar (mod mod), without requiring gcd(bas, mod) = 1.
inline std::optional<std::uint64_t> exbsgs(std::uint64_t bas, std::uint64_t tar,
                                           std::uint64_t mod) {
  return discrete_log(bas, tar, mod);
}

} // namespace noya

namespace noya {

namespace mod_kth_root_internal {

inline std::uint64_t prime_power(std::uint64_t pri, int exp) {
  std::uint64_t res = 1;
  while (exp-- > 0) {
    res *= pri;
  }
  return res;
}

/// Extract a p^exp-th root. The caller guarantees existence.
inline std::uint64_t extract_prime_power_root(std::uint64_t val,
                                              std::uint64_t p, int exp,
                                              std::uint64_t mod) {
  using factorize_internal::multiply_mod;
  using factorize_internal::power_mod;

  std::uint64_t cp = mod - 1;
  int vp = 0;
  while (cp % p == 0) {
    cp /= p;
    vp++;
  }

  const std::uint64_t pw = prime_power(p, exp);
  const std::uint64_t icp = discrete_log_internal::inverse_mod(cp % pw, pw);
  const std::uint64_t ie = std::uint64_t(
      ((static_cast<unsigned __int128>(pw - 1) * icp % pw) * cp + 1) / pw);
  std::uint64_t ans = power_mod(val, ie, mod);

  std::uint64_t nqr = 2;
  while (power_mod(nqr, (mod - 1) / p, mod) == 1) {
    nqr++;
  }
  nqr = power_mod(nqr, cp, mod);

  const int gs = int(std::sqrt(static_cast<long double>(vp - exp) * p)) + 1;
  const std::uint64_t tp = prime_power(p, vp - 1);
  const std::uint64_t ts =
      power_mod(nqr, factorize_internal::multiply_mod(gs, tp, mod - 1), mod);
  std::map<std::uint64_t, int> tbl;
  std::uint64_t cur = 1;
  for (int idx = 0; idx <= gs; idx++) {
    tbl[cur] = idx;
    cur = multiply_mod(cur, ts, mod);
  }

  const std::uint64_t is =
      discrete_log_internal::inverse_mod(power_mod(nqr, tp, mod), mod);
  for (int lev = exp; lev < vp; lev++) {
    std::uint64_t err = multiply_mod(
        discrete_log_internal::inverse_mod(power_mod(ans, pw, mod), mod), val,
        mod);
    std::uint64_t tar = power_mod(err, prime_power(p, vp - 1 - lev), mod);
    bool fou = false;
    for (int ste = 0; ste <= gs; ste++) {
      auto it = tbl.find(tar);
      if (it != tbl.end()) {
        const std::uint64_t dig =
            std::uint64_t(ste) + std::uint64_t(gs) * it->second;
        const std::uint64_t ce = factorize_internal::multiply_mod(
            dig, prime_power(p, lev - exp), mod - 1);
        ans = multiply_mod(ans, power_mod(nqr, ce, mod), mod);
        fou = true;
        break;
      }
      tar = multiply_mod(tar, is, mod);
    }
    assert(fou);
  }
  return ans;
}

} // namespace mod_kth_root_internal

/// @brief Return one x satisfying x^k = val modulo a prime, or nullopt;
/// avoids a full discrete logarithm by extracting the prime-power factors of
/// gcd(k, p - 1) independently.
inline std::optional<std::uint64_t>
mod_kth_root(std::uint64_t val, std::uint64_t exp, std::uint64_t pri) {
  assert(pri >= 2 && is_prime(pri));
  val %= pri;
  if (exp == 0) {
    return val == 1 ? std::optional<std::uint64_t>(1) : std::nullopt;
  }
  if (val == 0) {
    return 0;
  }
  if (pri == 2) {
    return val;
  }
  const std::uint64_t ord = pri - 1;
  exp %= ord;
  if (exp == 0) {
    return val == 1 ? std::optional<std::uint64_t>(1) : std::nullopt;
  }
  std::uint64_t div = std::gcd(exp, ord);
  if (factorize_internal::power_mod(val, ord / div, pri) != 1) {
    return std::nullopt;
  }

  const std::uint64_t ro = ord / div;
  if (ro > 1) {
    const std::uint64_t inv =
        discrete_log_internal::inverse_mod(exp / div % ro, ro);
    val = factorize_internal::power_mod(val, inv, pri);
  }

  for (std::uint64_t fct = 2; fct * fct <= div; fct++) {
    int cnt = 0;
    while (div % fct == 0) {
      div /= fct;
      cnt++;
    }
    if (cnt > 0) {
      val = mod_kth_root_internal::extract_prime_power_root(val, fct, cnt, pri);
    }
  }
  if (div > 1) {
    val = mod_kth_root_internal::extract_prime_power_root(val, div, 1, pri);
  }
  return val;
}

/// @brief Return the number of roots to x^k = val modulo a prime.
inline std::uint64_t mod_kth_root_count(std::uint64_t val, std::uint64_t exp,
                                        std::uint64_t pri) {
  assert(pri >= 2 && is_prime(pri));
  val %= pri;
  if (exp == 0) {
    return val == 1 ? pri : 0;
  }
  if (val == 0) {
    return 1;
  }
  auto rt = mod_kth_root(val, exp, pri);
  return rt ? std::gcd(exp, pri - 1) : 0;
}

} // namespace noya