mod_kth_root.hpp¶
Return one x satisfying x^k = value modulo a prime, or nullopt; avoids a full discrete logarithm by extracting the prime-power factors of gcd(k, p - 1) independently.
Verified by kth_root_mod.
\[
\displaystyle x^k\equiv a\pmod p
\]
Implementation¶
#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 prime, int exponent) {
std::uint64_t result = 1;
while (exponent-- > 0) {
result *= prime;
}
return result;
}
/// Extract a prime_power^exponent-th root. The caller guarantees existence.
inline std::uint64_t extract_prime_power_root(std::uint64_t value,
std::uint64_t prime_power_base,
int exponent,
std::uint64_t modulus) {
using factorize_internal::multiply_mod;
using factorize_internal::power_mod;
std::uint64_t coprime_part = modulus - 1;
int valuation = 0;
while (coprime_part % prime_power_base == 0) {
coprime_part /= prime_power_base;
valuation++;
}
const std::uint64_t power = prime_power(prime_power_base, exponent);
const std::uint64_t inverse_coprime =
discrete_log_internal::inverse_mod(coprime_part % power, power);
const std::uint64_t initial_exponent = std::uint64_t(
((static_cast<unsigned __int128>(power - 1) * inverse_coprime % power) *
coprime_part +
1) /
power);
std::uint64_t answer = power_mod(value, initial_exponent, modulus);
std::uint64_t non_residue = 2;
while (power_mod(non_residue, (modulus - 1) / prime_power_base,
modulus) == 1) {
non_residue++;
}
non_residue = power_mod(non_residue, coprime_part, modulus);
const int giant_step = int(std::sqrt(
static_cast<long double>(valuation - exponent) *
prime_power_base)) +
1;
const std::uint64_t top_power =
prime_power(prime_power_base, valuation - 1);
const std::uint64_t table_step = power_mod(
non_residue,
factorize_internal::multiply_mod(giant_step, top_power, modulus - 1),
modulus);
std::map<std::uint64_t, int> table;
std::uint64_t current = 1;
for (int index = 0; index <= giant_step; index++) {
table[current] = index;
current = multiply_mod(current, table_step, modulus);
}
const std::uint64_t inverse_small_step = discrete_log_internal::inverse_mod(
power_mod(non_residue, top_power, modulus), modulus);
for (int level = exponent; level < valuation; level++) {
std::uint64_t error = multiply_mod(
discrete_log_internal::inverse_mod(power_mod(answer, power, modulus),
modulus),
value, modulus);
std::uint64_t target = power_mod(
error, prime_power(prime_power_base, valuation - 1 - level), modulus);
bool found = false;
for (int step = 0; step <= giant_step; step++) {
auto iterator = table.find(target);
if (iterator != table.end()) {
const std::uint64_t digit =
std::uint64_t(step) + std::uint64_t(giant_step) * iterator->second;
const std::uint64_t correction_exponent =
factorize_internal::multiply_mod(
digit, prime_power(prime_power_base, level - exponent),
modulus - 1);
answer = multiply_mod(
answer, power_mod(non_residue, correction_exponent, modulus),
modulus);
found = true;
break;
}
target = multiply_mod(target, inverse_small_step, modulus);
}
assert(found);
}
return answer;
}
} // namespace mod_kth_root_internal
/// @brief Return one x satisfying x^k = value 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 value, std::uint64_t exponent,
std::uint64_t prime) {
assert(prime >= 2 && is_prime(prime));
value %= prime;
if (exponent == 0) {
return value == 1 ? std::optional<std::uint64_t>(1) : std::nullopt;
}
if (value == 0) {
return 0;
}
if (prime == 2) {
return value;
}
const std::uint64_t order = prime - 1;
exponent %= order;
if (exponent == 0) {
return value == 1 ? std::optional<std::uint64_t>(1) : std::nullopt;
}
std::uint64_t divisor = std::gcd(exponent, order);
if (factorize_internal::power_mod(value, order / divisor, prime) != 1) {
return std::nullopt;
}
const std::uint64_t reduced_order = order / divisor;
if (reduced_order > 1) {
const std::uint64_t inverse = discrete_log_internal::inverse_mod(
exponent / divisor % reduced_order, reduced_order);
value = factorize_internal::power_mod(value, inverse, prime);
}
for (std::uint64_t factor = 2; factor * factor <= divisor; factor++) {
int multiplicity = 0;
while (divisor % factor == 0) {
divisor /= factor;
multiplicity++;
}
if (multiplicity > 0) {
value = mod_kth_root_internal::extract_prime_power_root(
value, factor, multiplicity, prime);
}
}
if (divisor > 1) {
value = mod_kth_root_internal::extract_prime_power_root(value, divisor, 1,
prime);
}
return value;
}
/// @brief Return the number of roots to x^k = value modulo a prime.
inline std::uint64_t mod_kth_root_count(std::uint64_t value,
std::uint64_t exponent,
std::uint64_t prime) {
assert(prime >= 2 && is_prime(prime));
value %= prime;
if (exponent == 0) {
return value == 1 ? prime : 0;
}
if (value == 0) {
return 1;
}
auto root = mod_kth_root(value, exponent, prime);
return root ? std::gcd(exponent, prime - 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 exponent, u64 mod) {
u64 result = 1;
while (exponent > 0) {
if (exponent & 1) {
result = multiply_mod(result, a, mod);
}
a = multiply_mod(a, a, mod);
exponent >>= 1;
}
return result;
}
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 shift = __builtin_ctzll(n - 1);
u64 odd = (n - 1) >> shift;
for (u64 base :
std::array<u64, 7>{2, 325, 9375, 28178, 450775, 9780504, 1795265022}) {
if (base % n == 0) {
continue;
}
u64 value = power_mod(base % n, odd, n);
if (value == 1 || value == n - 1) {
continue;
}
bool composite = true;
for (int i = 1; i < shift; i++) {
value = multiply_mod(value, value, n);
if (value == n - 1) {
composite = false;
break;
}
}
if (composite) {
return false;
}
}
return true;
}
inline u64 splitmix64(u64 &state) {
u64 z = (state += 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 state = 0x123456789abcdef0ULL;
while (true) {
u64 y = splitmix64(state) % (n - 1) + 1;
u64 c = splitmix64(state) % (n - 1) + 1;
constexpr u64 block = 128;
u64 g = 1;
u64 r = 1;
u64 q = 1;
u64 x = 0;
u64 saved_y = 0;
auto next = [&](u64 value) {
return u64((u128(multiply_mod(value, value, n)) + c) % n);
};
while (g == 1) {
x = y;
for (u64 i = 0; i < r; i++) {
y = next(y);
}
for (u64 offset = 0; offset < r && g == 1; offset += block) {
saved_y = y;
for (u64 i = 0; i < std::min(block, r - offset); i++) {
y = next(y);
u64 difference = x > y ? x - y : y - x;
q = multiply_mod(q, difference, n);
}
g = std::gcd(q, n);
}
r <<= 1;
}
if (g == n) {
do {
saved_y = next(saved_y);
u64 difference = x > saved_y ? x - saved_y : saved_y - x;
g = std::gcd(difference, n);
} while (g == 1);
}
if (g != n) {
return g;
}
}
}
inline void collect_factors(u64 n, std::vector<u64> &result) {
if (n == 1) {
return;
}
if (miller_rabin(n)) {
result.push_back(n);
return;
}
u64 factor = pollard_rho(n);
collect_factors(factor, result);
collect_factors(n / factor, result);
}
} // 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> result;
factorize_internal::collect_factors(n, result);
std::sort(result.begin(), result.end());
return result;
}
/// @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>> result;
for (std::uint64_t p : prime_factors(n)) {
if (result.empty() || result.back().first != p) {
result.emplace_back(p, 1);
} else {
result.back().second++;
}
}
return result;
}
} // namespace noya
namespace noya {
namespace discrete_log_internal {
using u64 = std::uint64_t;
using i128 = __int128;
inline u64 inverse_mod(u64 value, u64 modulus) {
assert(modulus >= 2 && std::gcd(value, modulus) == 1);
i128 old_coefficient = 1;
i128 coefficient = 0;
u64 old_remainder = value;
u64 remainder = modulus;
while (remainder != 0) {
u64 quotient = old_remainder / remainder;
u64 next_remainder = old_remainder - quotient * remainder;
old_remainder = remainder;
remainder = next_remainder;
i128 next_coefficient = old_coefficient - i128(quotient) * coefficient;
old_coefficient = coefficient;
coefficient = next_coefficient;
}
old_coefficient %= i128(modulus);
if (old_coefficient < 0) {
old_coefficient += modulus;
}
return u64(old_coefficient);
}
inline u64 ceil_sqrt(u64 value) {
u64 result = u64(std::sqrt(static_cast<long double>(value)));
while (result != 0 && result > value / result) {
result--;
}
while (result == 0 || result <= value / result) {
if (result * result == value) {
return result;
}
result++;
}
return result;
}
inline std::optional<u64> coprime_bsgs(u64 base, u64 target, u64 modulus) {
using factorize_internal::multiply_mod;
using factorize_internal::power_mod;
assert(modulus >= 2 && std::gcd(base, modulus) == 1);
u64 block = ceil_sqrt(modulus);
std::unordered_map<u64, u64> baby_step;
baby_step.reserve(std::size_t(block * 2 + 1));
u64 value = 1;
for (u64 exponent = 0; exponent < block; exponent++) {
baby_step.emplace(value, exponent);
value = multiply_mod(value, base, modulus);
}
u64 inverse_block = power_mod(inverse_mod(base, modulus), block, modulus);
value = target;
for (u64 giant = 0; giant <= block; giant++) {
auto it = baby_step.find(value);
if (it != baby_step.end()) {
return giant * block + it->second;
}
value = multiply_mod(value, inverse_block, modulus);
}
return std::nullopt;
}
} // namespace discrete_log_internal
/// @brief Find the smallest x >= 0 satisfying base^x = target (mod modulus),
/// or nullopt if no solution exists.
inline std::optional<std::uint64_t>
discrete_log(std::uint64_t base, std::uint64_t target, std::uint64_t modulus) {
using discrete_log_internal::coprime_bsgs;
using factorize_internal::multiply_mod;
assert(modulus >= 1);
if (modulus == 1) {
return 0;
}
base %= modulus;
target %= modulus;
if (target == 1) {
return 0;
}
std::uint64_t removed = 0;
std::uint64_t accumulated = 1;
while (true) {
std::uint64_t divisor = std::gcd(base, modulus);
if (divisor == 1) {
break;
}
if (target == accumulated) {
return removed;
}
if (target % divisor != 0) {
return std::nullopt;
}
target /= divisor;
modulus /= divisor;
accumulated = multiply_mod(accumulated, base / divisor, modulus);
removed++;
}
if (target == accumulated) {
return removed;
}
std::uint64_t normalized = multiply_mod(
target, discrete_log_internal::inverse_mod(accumulated, modulus),
modulus);
auto result = coprime_bsgs(base % modulus, normalized, modulus);
if (!result) {
return std::nullopt;
}
return *result + removed;
}
} // namespace noya
namespace noya {
namespace mod_kth_root_internal {
inline std::uint64_t prime_power(std::uint64_t prime, int exponent) {
std::uint64_t result = 1;
while (exponent-- > 0) {
result *= prime;
}
return result;
}
/// Extract a prime_power^exponent-th root. The caller guarantees existence.
inline std::uint64_t extract_prime_power_root(std::uint64_t value,
std::uint64_t prime_power_base,
int exponent,
std::uint64_t modulus) {
using factorize_internal::multiply_mod;
using factorize_internal::power_mod;
std::uint64_t coprime_part = modulus - 1;
int valuation = 0;
while (coprime_part % prime_power_base == 0) {
coprime_part /= prime_power_base;
valuation++;
}
const std::uint64_t power = prime_power(prime_power_base, exponent);
const std::uint64_t inverse_coprime =
discrete_log_internal::inverse_mod(coprime_part % power, power);
const std::uint64_t initial_exponent = std::uint64_t(
((static_cast<unsigned __int128>(power - 1) * inverse_coprime % power) *
coprime_part +
1) /
power);
std::uint64_t answer = power_mod(value, initial_exponent, modulus);
std::uint64_t non_residue = 2;
while (power_mod(non_residue, (modulus - 1) / prime_power_base,
modulus) == 1) {
non_residue++;
}
non_residue = power_mod(non_residue, coprime_part, modulus);
const int giant_step = int(std::sqrt(
static_cast<long double>(valuation - exponent) *
prime_power_base)) +
1;
const std::uint64_t top_power =
prime_power(prime_power_base, valuation - 1);
const std::uint64_t table_step = power_mod(
non_residue,
factorize_internal::multiply_mod(giant_step, top_power, modulus - 1),
modulus);
std::map<std::uint64_t, int> table;
std::uint64_t current = 1;
for (int index = 0; index <= giant_step; index++) {
table[current] = index;
current = multiply_mod(current, table_step, modulus);
}
const std::uint64_t inverse_small_step = discrete_log_internal::inverse_mod(
power_mod(non_residue, top_power, modulus), modulus);
for (int level = exponent; level < valuation; level++) {
std::uint64_t error = multiply_mod(
discrete_log_internal::inverse_mod(power_mod(answer, power, modulus),
modulus),
value, modulus);
std::uint64_t target = power_mod(
error, prime_power(prime_power_base, valuation - 1 - level), modulus);
bool found = false;
for (int step = 0; step <= giant_step; step++) {
auto iterator = table.find(target);
if (iterator != table.end()) {
const std::uint64_t digit =
std::uint64_t(step) + std::uint64_t(giant_step) * iterator->second;
const std::uint64_t correction_exponent =
factorize_internal::multiply_mod(
digit, prime_power(prime_power_base, level - exponent),
modulus - 1);
answer = multiply_mod(
answer, power_mod(non_residue, correction_exponent, modulus),
modulus);
found = true;
break;
}
target = multiply_mod(target, inverse_small_step, modulus);
}
assert(found);
}
return answer;
}
} // namespace mod_kth_root_internal
/// @brief Return one x satisfying x^k = value 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 value, std::uint64_t exponent,
std::uint64_t prime) {
assert(prime >= 2 && is_prime(prime));
value %= prime;
if (exponent == 0) {
return value == 1 ? std::optional<std::uint64_t>(1) : std::nullopt;
}
if (value == 0) {
return 0;
}
if (prime == 2) {
return value;
}
const std::uint64_t order = prime - 1;
exponent %= order;
if (exponent == 0) {
return value == 1 ? std::optional<std::uint64_t>(1) : std::nullopt;
}
std::uint64_t divisor = std::gcd(exponent, order);
if (factorize_internal::power_mod(value, order / divisor, prime) != 1) {
return std::nullopt;
}
const std::uint64_t reduced_order = order / divisor;
if (reduced_order > 1) {
const std::uint64_t inverse = discrete_log_internal::inverse_mod(
exponent / divisor % reduced_order, reduced_order);
value = factorize_internal::power_mod(value, inverse, prime);
}
for (std::uint64_t factor = 2; factor * factor <= divisor; factor++) {
int multiplicity = 0;
while (divisor % factor == 0) {
divisor /= factor;
multiplicity++;
}
if (multiplicity > 0) {
value = mod_kth_root_internal::extract_prime_power_root(
value, factor, multiplicity, prime);
}
}
if (divisor > 1) {
value = mod_kth_root_internal::extract_prime_power_root(value, divisor, 1,
prime);
}
return value;
}
/// @brief Return the number of roots to x^k = value modulo a prime.
inline std::uint64_t mod_kth_root_count(std::uint64_t value,
std::uint64_t exponent,
std::uint64_t prime) {
assert(prime >= 2 && is_prime(prime));
value %= prime;
if (exponent == 0) {
return value == 1 ? prime : 0;
}
if (value == 0) {
return 1;
}
auto root = mod_kth_root(value, exponent, prime);
return root ? std::gcd(exponent, prime - 1) : 0;
}
} // namespace noya