Skip to content

fixed_discrete_log.hpp

SECTIONMath INCLUDEnoya/fixed_discrete_log.hpp

Preprocess discrete logarithms to one primitive root modulo a prime. Farey approximation writes every nonzero x as xq = +/-t (mod p) with q <= p^(1/3) and |t| <= p^(2/3). Logs of the small integers are built from a batched baby-step/giant-step pass for primes, multiplicativity for small composites, and p = floor(p/i)i + (p mod i) for the remaining interval. A query then needs only the precomputed logs of q and |t|.

Verified by discrete_logarithm_fixed_mod.

\[ \displaystyle a^x\equiv b\pmod p,\quad x=\min\{t\ge 0:a^t\equiv b\pmod p\} \]

Implementation

View on GitHub

#ifndef NOYA_FIXED_DISCRETE_LOG_HPP
#define NOYA_FIXED_DISCRETE_LOG_HPP 1

/// @complexity Time: O(p^(2/3) + sqrt(p*pi(sqrt p)) log p)
/// preprocessing and O(1) per logarithm query for prime p.
/// Space: O(p^(2/3) + sqrt(p*pi(sqrt p))).

#include "noya/factorize.hpp"

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <limits>
#include <utility>
#include <vector>

namespace noya {

/// @brief Preprocess discrete logarithms to one primitive root modulo a prime.
/// Farey approximation writes every nonzero x as x*q = +/-t (mod p) with
/// q <= p^(1/3) and |t| <= p^(2/3). Logs of the small integers are built from
/// a batched baby-step/giant-step pass for primes, multiplicativity for small
/// composites, and p = floor(p/i)*i + (p mod i) for the remaining interval.
/// A query then needs only the precomputed logs of q and |t|.
class fixed_discrete_log_table {
  using u32 = std::uint32_t;
  using u64 = std::uint64_t;
  using fraction = std::pair<u32, u32>;

public:
  fixed_discrete_log_table(u32 prime, u32 primitive_root)
      : prime_(prime), generator_(primitive_root), order_(prime - 1) {
    assert(prime >= 2 && is_prime(prime));
    assert(generator_ > 0 && generator_ < prime_);
    if (prime_ == 2) {
      direct_log_ = {std::numeric_limits<u32>::max(), 0};
    } else if (prime_ <= direct_limit) {
      build_direct();
    } else {
      build_fast();
    }
  }

  u32 logarithm(u32 value) const {
    assert(value > 0 && value < prime_);
    if (!direct_log_.empty()) {
      return direct_log_[value];
    }

    u32 index = u32(u64(value) * farey_scale_ / prime_);
    auto [numerator, denominator] = predecessor_[index];
    std::int64_t difference = std::int64_t(u64(value) * denominator) -
                              std::int64_t(u64(prime_) * numerator);
    if (std::uint64_t(std::abs(difference)) > farey_scale_) {
      numerator = successor_[index].first;
      denominator = successor_[index].second;
      difference = std::int64_t(u64(value) * denominator) -
                   std::int64_t(u64(prime_) * numerator);
    }
    assert(difference != 0 &&
           std::uint64_t(std::abs(difference)) <= farey_scale_);
    u32 numerator_log = small_log_[std::size_t(std::abs(difference))];
    if (difference < 0) {
      numerator_log = add_exponents(numerator_log, order_ / 2);
    }
    return subtract_exponents(numerator_log, small_log_[denominator]);
  }

private:
  static constexpr u32 direct_limit = 1'000'000;

  u32 prime_;
  u32 generator_;
  u32 order_;
  u32 farey_scale_ = 0;
  std::vector<u32> direct_log_;
  std::vector<u32> small_log_;
  std::vector<fraction> predecessor_;
  std::vector<fraction> successor_;

  u32 multiply(u32 first, u32 second) const {
    return u32(u64(first) * second % prime_);
  }

  u32 power(u32 value, u64 exponent) const {
    return u32(factorize_internal::power_mod(value, exponent, prime_));
  }

  u32 add_exponents(u32 first, u32 second) const {
    u32 result = first + second;
    return result >= order_ ? result - order_ : result;
  }

  u32 subtract_exponents(u32 first, u32 second) const {
    return first >= second ? first - second : first + order_ - second;
  }

  void build_direct() {
    direct_log_.assign(prime_, std::numeric_limits<u32>::max());
    u32 value = 1;
    for (u32 exponent = 0; exponent < order_; exponent++) {
      assert(direct_log_[value] == std::numeric_limits<u32>::max());
      direct_log_[value] = exponent;
      value = multiply(value, generator_);
    }
    assert(value == 1);
  }

  std::vector<u32> smallest_prime_factors(u32 limit,
                                           std::vector<u32> &primes) const {
    std::vector<u32> smallest(limit + 1);
    for (u32 value = 2; value <= limit; value++) {
      if (smallest[value] == 0) {
        smallest[value] = value;
        primes.push_back(value);
      }
      for (u32 prime : primes) {
        if (prime > smallest[value] || u64(value) * prime > limit) {
          break;
        }
        smallest[value * prime] = prime;
      }
    }
    return smallest;
  }

  std::vector<u32> batch_prime_logs(const std::vector<u32> &targets) const {
    if (targets.empty()) {
      return {};
    }
    u32 block = u32(std::sqrt(
        static_cast<long double>(prime_) / targets.size())) + 2;
    u32 giant_count = prime_ / block + 3;
    u32 giant_step = power(generator_, block);
    std::vector<std::pair<u32, u32>> giants;
    giants.reserve(giant_count);
    u32 value = giant_step;
    for (u32 x = 1; x <= giant_count; x++) {
      giants.emplace_back(value, x);
      value = multiply(value, giant_step);
    }
    std::sort(giants.begin(), giants.end());

    std::vector<u32> answer(targets.size(),
                            std::numeric_limits<u32>::max());
    u32 baby = 1;
    for (u32 y = 0; y < block; y++) {
      for (std::size_t index = 0; index < targets.size(); index++) {
        u32 wanted = multiply(targets[index], baby);
        auto iterator = std::lower_bound(
            giants.begin(), giants.end(), std::pair<u32, u32>{wanted, 0});
        if (iterator != giants.end() && iterator->first == wanted) {
          u64 candidate = u64(iterator->second) * block - y;
          if (candidate < answer[index]) {
            answer[index] = u32(candidate);
          }
        }
      }
      baby = multiply(baby, generator_);
    }
    for (u32 exponent : answer) {
      assert(exponent < order_);
    }
    return answer;
  }

  void build_fast() {
    u32 farey_bound = 1;
    while (u64(farey_bound) * farey_bound * farey_bound <= prime_) {
      farey_bound *= 2;
    }
    farey_scale_ = farey_bound * farey_bound;

    std::vector<fraction> exact(farey_scale_ + 1);
    for (u32 numerator = 0; numerator <= farey_bound; numerator++) {
      u32 first_denominator = numerator == 1 ? 1 : numerator + 1;
      for (u32 denominator = first_denominator;
           denominator <= farey_bound; denominator++) {
        u32 index = u32(u64(numerator) * farey_scale_ / denominator);
        if (exact[index].second == 0) {
          exact[index] = {numerator, denominator};
        }
      }
    }
    predecessor_.resize(farey_scale_ + 1);
    fraction current{0, 1};
    for (u32 index = 0; index <= farey_scale_; index++) {
      if (exact[index].second != 0) {
        current = exact[index];
      }
      predecessor_[index] = current;
    }
    successor_.resize(farey_scale_ + 1);
    current = {1, 1};
    for (u32 index = farey_scale_;; index--) {
      if (exact[index].second != 0) {
        current = exact[index];
      }
      successor_[index] = current;
      if (index == 0) {
        break;
      }
    }

    u32 square_root = u32(std::sqrt(static_cast<long double>(prime_)));
    while (u64(square_root) * square_root > prime_) {
      square_root--;
    }
    while (u64(square_root + 1) * (square_root + 1) <= prime_) {
      square_root++;
    }
    std::vector<u32> primes;
    std::vector<u32> smallest = smallest_prime_factors(square_root, primes);
    std::vector<u32> prime_logs = batch_prime_logs(primes);

    small_log_.assign(farey_scale_ + 1, 0);
    for (std::size_t index = 0; index < primes.size(); index++) {
      small_log_[primes[index]] = prime_logs[index];
    }
    for (u32 value = 2; value <= square_root; value++) {
      if (smallest[value] != value) {
        small_log_[value] = add_exponents(
            small_log_[smallest[value]], small_log_[value / smallest[value]]);
      }
    }
    for (u32 value = square_root + 1; value <= farey_scale_; value++) {
      u32 quotient = prime_ / value;
      u32 remainder = prime_ % value;
      small_log_[value] = subtract_exponents(
          add_exponents(order_ / 2, small_log_[remainder]),
          small_log_[quotient]);
    }
  }
};

} // namespace noya

#endif // NOYA_FIXED_DISCRETE_LOG_HPP
#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <limits>
#include <numeric>
#include <utility>
#include <vector>

/// @complexity Time: O(p^(2/3) + sqrt(p*pi(sqrt p)) log p)
/// preprocessing and O(1) per logarithm query for prime p.
/// Space: O(p^(2/3) + sqrt(p*pi(sqrt p))).

/// @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 {

/// @brief Preprocess discrete logarithms to one primitive root modulo a prime.
/// Farey approximation writes every nonzero x as x*q = +/-t (mod p) with
/// q <= p^(1/3) and |t| <= p^(2/3). Logs of the small integers are built from
/// a batched baby-step/giant-step pass for primes, multiplicativity for small
/// composites, and p = floor(p/i)*i + (p mod i) for the remaining interval.
/// A query then needs only the precomputed logs of q and |t|.
class fixed_discrete_log_table {
  using u32 = std::uint32_t;
  using u64 = std::uint64_t;
  using fraction = std::pair<u32, u32>;

public:
  fixed_discrete_log_table(u32 prime, u32 primitive_root)
      : prime_(prime), generator_(primitive_root), order_(prime - 1) {
    assert(prime >= 2 && is_prime(prime));
    assert(generator_ > 0 && generator_ < prime_);
    if (prime_ == 2) {
      direct_log_ = {std::numeric_limits<u32>::max(), 0};
    } else if (prime_ <= direct_limit) {
      build_direct();
    } else {
      build_fast();
    }
  }

  u32 logarithm(u32 value) const {
    assert(value > 0 && value < prime_);
    if (!direct_log_.empty()) {
      return direct_log_[value];
    }

    u32 index = u32(u64(value) * farey_scale_ / prime_);
    auto [numerator, denominator] = predecessor_[index];
    std::int64_t difference = std::int64_t(u64(value) * denominator) -
                              std::int64_t(u64(prime_) * numerator);
    if (std::uint64_t(std::abs(difference)) > farey_scale_) {
      numerator = successor_[index].first;
      denominator = successor_[index].second;
      difference = std::int64_t(u64(value) * denominator) -
                   std::int64_t(u64(prime_) * numerator);
    }
    assert(difference != 0 &&
           std::uint64_t(std::abs(difference)) <= farey_scale_);
    u32 numerator_log = small_log_[std::size_t(std::abs(difference))];
    if (difference < 0) {
      numerator_log = add_exponents(numerator_log, order_ / 2);
    }
    return subtract_exponents(numerator_log, small_log_[denominator]);
  }

private:
  static constexpr u32 direct_limit = 1'000'000;

  u32 prime_;
  u32 generator_;
  u32 order_;
  u32 farey_scale_ = 0;
  std::vector<u32> direct_log_;
  std::vector<u32> small_log_;
  std::vector<fraction> predecessor_;
  std::vector<fraction> successor_;

  u32 multiply(u32 first, u32 second) const {
    return u32(u64(first) * second % prime_);
  }

  u32 power(u32 value, u64 exponent) const {
    return u32(factorize_internal::power_mod(value, exponent, prime_));
  }

  u32 add_exponents(u32 first, u32 second) const {
    u32 result = first + second;
    return result >= order_ ? result - order_ : result;
  }

  u32 subtract_exponents(u32 first, u32 second) const {
    return first >= second ? first - second : first + order_ - second;
  }

  void build_direct() {
    direct_log_.assign(prime_, std::numeric_limits<u32>::max());
    u32 value = 1;
    for (u32 exponent = 0; exponent < order_; exponent++) {
      assert(direct_log_[value] == std::numeric_limits<u32>::max());
      direct_log_[value] = exponent;
      value = multiply(value, generator_);
    }
    assert(value == 1);
  }

  std::vector<u32> smallest_prime_factors(u32 limit,
                                           std::vector<u32> &primes) const {
    std::vector<u32> smallest(limit + 1);
    for (u32 value = 2; value <= limit; value++) {
      if (smallest[value] == 0) {
        smallest[value] = value;
        primes.push_back(value);
      }
      for (u32 prime : primes) {
        if (prime > smallest[value] || u64(value) * prime > limit) {
          break;
        }
        smallest[value * prime] = prime;
      }
    }
    return smallest;
  }

  std::vector<u32> batch_prime_logs(const std::vector<u32> &targets) const {
    if (targets.empty()) {
      return {};
    }
    u32 block = u32(std::sqrt(
        static_cast<long double>(prime_) / targets.size())) + 2;
    u32 giant_count = prime_ / block + 3;
    u32 giant_step = power(generator_, block);
    std::vector<std::pair<u32, u32>> giants;
    giants.reserve(giant_count);
    u32 value = giant_step;
    for (u32 x = 1; x <= giant_count; x++) {
      giants.emplace_back(value, x);
      value = multiply(value, giant_step);
    }
    std::sort(giants.begin(), giants.end());

    std::vector<u32> answer(targets.size(),
                            std::numeric_limits<u32>::max());
    u32 baby = 1;
    for (u32 y = 0; y < block; y++) {
      for (std::size_t index = 0; index < targets.size(); index++) {
        u32 wanted = multiply(targets[index], baby);
        auto iterator = std::lower_bound(
            giants.begin(), giants.end(), std::pair<u32, u32>{wanted, 0});
        if (iterator != giants.end() && iterator->first == wanted) {
          u64 candidate = u64(iterator->second) * block - y;
          if (candidate < answer[index]) {
            answer[index] = u32(candidate);
          }
        }
      }
      baby = multiply(baby, generator_);
    }
    for (u32 exponent : answer) {
      assert(exponent < order_);
    }
    return answer;
  }

  void build_fast() {
    u32 farey_bound = 1;
    while (u64(farey_bound) * farey_bound * farey_bound <= prime_) {
      farey_bound *= 2;
    }
    farey_scale_ = farey_bound * farey_bound;

    std::vector<fraction> exact(farey_scale_ + 1);
    for (u32 numerator = 0; numerator <= farey_bound; numerator++) {
      u32 first_denominator = numerator == 1 ? 1 : numerator + 1;
      for (u32 denominator = first_denominator;
           denominator <= farey_bound; denominator++) {
        u32 index = u32(u64(numerator) * farey_scale_ / denominator);
        if (exact[index].second == 0) {
          exact[index] = {numerator, denominator};
        }
      }
    }
    predecessor_.resize(farey_scale_ + 1);
    fraction current{0, 1};
    for (u32 index = 0; index <= farey_scale_; index++) {
      if (exact[index].second != 0) {
        current = exact[index];
      }
      predecessor_[index] = current;
    }
    successor_.resize(farey_scale_ + 1);
    current = {1, 1};
    for (u32 index = farey_scale_;; index--) {
      if (exact[index].second != 0) {
        current = exact[index];
      }
      successor_[index] = current;
      if (index == 0) {
        break;
      }
    }

    u32 square_root = u32(std::sqrt(static_cast<long double>(prime_)));
    while (u64(square_root) * square_root > prime_) {
      square_root--;
    }
    while (u64(square_root + 1) * (square_root + 1) <= prime_) {
      square_root++;
    }
    std::vector<u32> primes;
    std::vector<u32> smallest = smallest_prime_factors(square_root, primes);
    std::vector<u32> prime_logs = batch_prime_logs(primes);

    small_log_.assign(farey_scale_ + 1, 0);
    for (std::size_t index = 0; index < primes.size(); index++) {
      small_log_[primes[index]] = prime_logs[index];
    }
    for (u32 value = 2; value <= square_root; value++) {
      if (smallest[value] != value) {
        small_log_[value] = add_exponents(
            small_log_[smallest[value]], small_log_[value / smallest[value]]);
      }
    }
    for (u32 value = square_root + 1; value <= farey_scale_; value++) {
      u32 quotient = prime_ / value;
      u32 remainder = prime_ % value;
      small_log_[value] = subtract_exponents(
          add_exponents(order_ / 2, small_log_[remainder]),
          small_log_[quotient]);
    }
  }
};

} // namespace noya