binomial_mod.hpp

SECTIONMath INCLUDEnoya/binomial_mod.hpp

Preprocess one fixed composite modulus for many binomial queries. Each prime-power factor stores prefix products with multiples of its prime removed. Splitting n! into complete residue blocks and recursively stripping one prime from every multiple yields its unit part and p-adic valuation in logarithmic time. The prime-power answers are joined by CRT.

Verified by binomial_coefficient.

\[ \displaystyle \binom{n}{k} \bmod m \]

Implementation

View on GitHub

SourceExpanded

#ifndef NOYA_BINOMIAL_MOD_HPP
#define NOYA_BINOMIAL_MOD_HPP 1

/// @complexity Time: O(m) preprocessing and O(sum_{p|m} log_p n) per table
/// query; O(min(k,n-k) omega(m)) for the one-shot function.
/// Space: O(m) for the table and O(omega(m)) for the one-shot function.

#include "noya/factorize.hpp"

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <vector>

namespace noya {
namespace binomial_mod_internal {

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

inline u64 multiply(u64 a, u64 b, u64 modulus) {
  return u64(u128(a) * b % modulus);
}

inline u64 power(u64 value, u64 exponent, u64 modulus) {
  u64 result = 1 % modulus;
  while (exponent > 0) {
    if (exponent & 1) {
      result = multiply(result, value, modulus);
    }
    value = multiply(value, value, modulus);
    exponent >>= 1;
  }
  return result;
}

inline u64 inverse(u64 value, u64 modulus) {
  i128 old_r = value;
  i128 r = modulus;
  i128 old_s = 1;
  i128 s = 0;
  while (r != 0) {
    i128 quotient = old_r / r;
    i128 next_r = old_r - quotient * r;
    old_r = r;
    r = next_r;
    i128 next_s = old_s - quotient * s;
    old_s = s;
    s = next_s;
  }
  assert(old_r == 1);
  old_s %= i128(modulus);
  if (old_s < 0) {
    old_s += modulus;
  }
  return u64(old_s);
}

inline u64 prime_power_residue(u64 n, u64 k, u64 prime, int exponent,
                               u64 modulus) {
  std::int64_t valuation = 0;
  u64 unit = 1 % modulus;
  for (u64 index = 1; index <= k; index++) {
    u64 numerator = n - k + index;
    u64 denominator = index;
    while (numerator % prime == 0) {
      numerator /= prime;
      valuation++;
    }
    while (denominator % prime == 0) {
      denominator /= prime;
      valuation--;
    }
    unit = multiply(unit, numerator % modulus, modulus);
    unit = multiply(unit, inverse(denominator % modulus, modulus), modulus);
  }
  assert(valuation >= 0);
  if (valuation >= exponent) {
    return 0;
  }
  return multiply(unit, power(prime, valuation, modulus), modulus);
}

struct prime_power_binomial {
  u64 prime = 0;
  int exponent = 0;
  u64 modulus = 1;
  std::vector<std::uint32_t> unit_prefix;

  prime_power_binomial(u64 prime_, int exponent_)
      : prime(prime_), exponent(exponent_) {
    for (int i = 0; i < exponent; i++) {
      modulus *= prime;
    }
    unit_prefix.resize(std::size_t(modulus) + 1);
    unit_prefix[0] = 1;
    for (u64 i = 1; i <= modulus; i++) {
      unit_prefix[i] = unit_prefix[i - 1];
      if (i % prime != 0) {
        unit_prefix[i] = std::uint32_t(
            multiply(unit_prefix[i], i, modulus));
      }
    }
  }

  u64 valuation_factorial(u64 n) const {
    u64 result = 0;
    while (n > 0) {
      n /= prime;
      result += n;
    }
    return result;
  }

  u64 unit_factorial(u64 n) const {
    u64 result = 1;
    while (n > 0) {
      if ((n / modulus) & 1) {
        result = multiply(result, unit_prefix[modulus], modulus);
      }
      result = multiply(result, unit_prefix[n % modulus], modulus);
      n /= prime;
    }
    return result;
  }

  u64 choose(u64 n, u64 k) const {
    if (k > n) {
      return 0;
    }
    u64 valuation = valuation_factorial(n) - valuation_factorial(k) -
                    valuation_factorial(n - k);
    if (valuation >= u64(exponent)) {
      return 0;
    }
    u64 unit = unit_factorial(n);
    unit = multiply(unit, inverse(unit_factorial(k), modulus), modulus);
    unit = multiply(unit, inverse(unit_factorial(n - k), modulus), modulus);
    return multiply(unit, power(prime, valuation, modulus), modulus);
  }
};

} // namespace binomial_mod_internal

/// @brief Preprocess one fixed composite modulus for many binomial queries.
/// Each prime-power factor stores prefix products with multiples of its prime
/// removed. Splitting n! into complete residue blocks and recursively stripping
/// one prime from every multiple yields its unit part and p-adic valuation in
/// logarithmic time. The prime-power answers are joined by CRT.
class binomial_mod_table {
public:
  explicit binomial_mod_table(std::uint32_t modulus) : modulus_(modulus) {
    assert(modulus >= 1);
    if (modulus == 1) {
      return;
    }
    for (auto [prime, exponent] : factorize(modulus)) {
      factors_.emplace_back(prime, exponent);
    }
  }

  std::uint32_t modulus() const { return modulus_; }

  std::uint32_t choose(std::uint64_t n, std::uint64_t k) const {
    using namespace binomial_mod_internal;
    if (modulus_ == 1 || k > n) {
      return 0;
    }
    u64 result = 0;
    u64 combined_modulus = 1;
    for (const auto &factor : factors_) {
      u64 residue = factor.choose(n, k);
      u64 current = result % factor.modulus;
      u64 difference =
          (residue + factor.modulus - current) % factor.modulus;
      u64 scale = multiply(
          difference,
          inverse(combined_modulus % factor.modulus, factor.modulus),
          factor.modulus);
      result += combined_modulus * scale;
      combined_modulus *= factor.modulus;
      result %= combined_modulus;
    }
    return std::uint32_t(result);
  }

private:
  std::uint32_t modulus_;
  std::vector<binomial_mod_internal::prime_power_binomial> factors_;
};

/// @brief Compute C(n,k) modulo an arbitrary 64-bit positive modulus in
/// O(min(k,n-k) times the number of prime factors of modulus).
inline std::uint64_t binomial_mod(std::uint64_t n, std::uint64_t k,
                                  std::uint64_t modulus) {
  using namespace binomial_mod_internal;
  assert(modulus >= 1);
  if (modulus == 1 || k > n) {
    return 0;
  }
  k = std::min(k, n - k);
  u64 result = 0;
  u64 combined_modulus = 1;
  for (auto [prime, exponent] : factorize(modulus)) {
    u64 prime_power = 1;
    for (int count = 0; count < exponent; count++) {
      prime_power *= prime;
    }
    u64 residue = prime_power_residue(n, k, prime, exponent, prime_power);
    u64 current = result % prime_power;
    u64 difference = u64((u128(residue) + prime_power - current) % prime_power);
    u64 scale = multiply(difference,
                         inverse(combined_modulus % prime_power, prime_power),
                         prime_power);
    result = u64(u128(result) + u128(combined_modulus) * scale);
    combined_modulus *= prime_power;
    result %= combined_modulus;
  }
  return result;
}

} // namespace noya

#endif // NOYA_BINOMIAL_MOD_HPP

#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <numeric>
#include <utility>
#include <vector>

/// @complexity Time: O(m) preprocessing and O(sum_{p|m} log_p n) per table
/// query; O(min(k,n-k) omega(m)) for the one-shot function.
/// Space: O(m) for the table and O(omega(m)) for the one-shot function.

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

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

inline u64 multiply(u64 a, u64 b, u64 modulus) {
  return u64(u128(a) * b % modulus);
}

inline u64 power(u64 value, u64 exponent, u64 modulus) {
  u64 result = 1 % modulus;
  while (exponent > 0) {
    if (exponent & 1) {
      result = multiply(result, value, modulus);
    }
    value = multiply(value, value, modulus);
    exponent >>= 1;
  }
  return result;
}

inline u64 inverse(u64 value, u64 modulus) {
  i128 old_r = value;
  i128 r = modulus;
  i128 old_s = 1;
  i128 s = 0;
  while (r != 0) {
    i128 quotient = old_r / r;
    i128 next_r = old_r - quotient * r;
    old_r = r;
    r = next_r;
    i128 next_s = old_s - quotient * s;
    old_s = s;
    s = next_s;
  }
  assert(old_r == 1);
  old_s %= i128(modulus);
  if (old_s < 0) {
    old_s += modulus;
  }
  return u64(old_s);
}

inline u64 prime_power_residue(u64 n, u64 k, u64 prime, int exponent,
                               u64 modulus) {
  std::int64_t valuation = 0;
  u64 unit = 1 % modulus;
  for (u64 index = 1; index <= k; index++) {
    u64 numerator = n - k + index;
    u64 denominator = index;
    while (numerator % prime == 0) {
      numerator /= prime;
      valuation++;
    }
    while (denominator % prime == 0) {
      denominator /= prime;
      valuation--;
    }
    unit = multiply(unit, numerator % modulus, modulus);
    unit = multiply(unit, inverse(denominator % modulus, modulus), modulus);
  }
  assert(valuation >= 0);
  if (valuation >= exponent) {
    return 0;
  }
  return multiply(unit, power(prime, valuation, modulus), modulus);
}

struct prime_power_binomial {
  u64 prime = 0;
  int exponent = 0;
  u64 modulus = 1;
  std::vector<std::uint32_t> unit_prefix;

  prime_power_binomial(u64 prime_, int exponent_)
      : prime(prime_), exponent(exponent_) {
    for (int i = 0; i < exponent; i++) {
      modulus *= prime;
    }
    unit_prefix.resize(std::size_t(modulus) + 1);
    unit_prefix[0] = 1;
    for (u64 i = 1; i <= modulus; i++) {
      unit_prefix[i] = unit_prefix[i - 1];
      if (i % prime != 0) {
        unit_prefix[i] = std::uint32_t(
            multiply(unit_prefix[i], i, modulus));
      }
    }
  }

  u64 valuation_factorial(u64 n) const {
    u64 result = 0;
    while (n > 0) {
      n /= prime;
      result += n;
    }
    return result;
  }

  u64 unit_factorial(u64 n) const {
    u64 result = 1;
    while (n > 0) {
      if ((n / modulus) & 1) {
        result = multiply(result, unit_prefix[modulus], modulus);
      }
      result = multiply(result, unit_prefix[n % modulus], modulus);
      n /= prime;
    }
    return result;
  }

  u64 choose(u64 n, u64 k) const {
    if (k > n) {
      return 0;
    }
    u64 valuation = valuation_factorial(n) - valuation_factorial(k) -
                    valuation_factorial(n - k);
    if (valuation >= u64(exponent)) {
      return 0;
    }
    u64 unit = unit_factorial(n);
    unit = multiply(unit, inverse(unit_factorial(k), modulus), modulus);
    unit = multiply(unit, inverse(unit_factorial(n - k), modulus), modulus);
    return multiply(unit, power(prime, valuation, modulus), modulus);
  }
};

} // namespace binomial_mod_internal

/// @brief Preprocess one fixed composite modulus for many binomial queries.
/// Each prime-power factor stores prefix products with multiples of its prime
/// removed. Splitting n! into complete residue blocks and recursively stripping
/// one prime from every multiple yields its unit part and p-adic valuation in
/// logarithmic time. The prime-power answers are joined by CRT.
class binomial_mod_table {
public:
  explicit binomial_mod_table(std::uint32_t modulus) : modulus_(modulus) {
    assert(modulus >= 1);
    if (modulus == 1) {
      return;
    }
    for (auto [prime, exponent] : factorize(modulus)) {
      factors_.emplace_back(prime, exponent);
    }
  }

  std::uint32_t modulus() const { return modulus_; }

  std::uint32_t choose(std::uint64_t n, std::uint64_t k) const {
    using namespace binomial_mod_internal;
    if (modulus_ == 1 || k > n) {
      return 0;
    }
    u64 result = 0;
    u64 combined_modulus = 1;
    for (const auto &factor : factors_) {
      u64 residue = factor.choose(n, k);
      u64 current = result % factor.modulus;
      u64 difference =
          (residue + factor.modulus - current) % factor.modulus;
      u64 scale = multiply(
          difference,
          inverse(combined_modulus % factor.modulus, factor.modulus),
          factor.modulus);
      result += combined_modulus * scale;
      combined_modulus *= factor.modulus;
      result %= combined_modulus;
    }
    return std::uint32_t(result);
  }

private:
  std::uint32_t modulus_;
  std::vector<binomial_mod_internal::prime_power_binomial> factors_;
};

/// @brief Compute C(n,k) modulo an arbitrary 64-bit positive modulus in
/// O(min(k,n-k) times the number of prime factors of modulus).
inline std::uint64_t binomial_mod(std::uint64_t n, std::uint64_t k,
                                  std::uint64_t modulus) {
  using namespace binomial_mod_internal;
  assert(modulus >= 1);
  if (modulus == 1 || k > n) {
    return 0;
  }
  k = std::min(k, n - k);
  u64 result = 0;
  u64 combined_modulus = 1;
  for (auto [prime, exponent] : factorize(modulus)) {
    u64 prime_power = 1;
    for (int count = 0; count < exponent; count++) {
      prime_power *= prime;
    }
    u64 residue = prime_power_residue(n, k, prime, exponent, prime_power);
    u64 current = result % prime_power;
    u64 difference = u64((u128(residue) + prime_power - current) % prime_power);
    u64 scale = multiply(difference,
                         inverse(combined_modulus % prime_power, prime_power),
                         prime_power);
    result = u64(u128(result) + u128(combined_modulus) * scale);
    combined_modulus *= prime_power;
    result %= combined_modulus;
  }
  return result;
}

} // namespace noya