q_binomial_prime.hpp¶
Answer Gaussian binomial coefficients at a fixed q modulo a prime. If d is the multiplicative order of q, q-Lucas gives [n choose k]_q = C(floor(n/d),floor(k/d)) [n mod d choose k mod d]_q. Ordinary factorials handle the first factor, while products of 1-q^i for i<d handle the second. q=0 and q=1 are treated by their direct limits.
Verified by q_binomial_coefficient_prime_mod.
\[
\displaystyle \binom{n}{k}_q
\]
Implementation¶
#ifndef NOYA_Q_BINOMIAL_PRIME_HPP
#define NOYA_Q_BINOMIAL_PRIME_HPP 1
/// @complexity Time: expected factorization time for p-1 plus
/// O(max(n/order) + min(order,max n) + T) for T queries.
/// Space: O(max(n/order) + min(order,max n)).
#include "noya/factorize.hpp"
#include "noya/prime_binomial_table.hpp"
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <utility>
#include <vector>
namespace noya {
/// @brief Answer Gaussian binomial coefficients at a fixed q modulo a prime.
/// If d is the multiplicative order of q, q-Lucas gives
/// [n choose k]_q = C(floor(n/d),floor(k/d)) [n mod d choose k mod d]_q.
/// Ordinary factorials handle the first factor, while products of 1-q^i for
/// i<d handle the second. q=0 and q=1 are treated by their direct limits.
class q_binomial_prime_table {
public:
q_binomial_prime_table() = default;
q_binomial_prime_table(
const std::vector<std::pair<int, int>> &queries, std::uint32_t prime,
std::uint32_t q) {
build(queries, prime, q);
}
void build(const std::vector<std::pair<int, int>> &queries,
std::uint32_t prime, std::uint32_t q) {
modulus_ = prime;
q_ = q % prime;
int maximum = 0;
for (auto [n, k] : queries) {
assert(n >= 0 && k >= 0 && std::uint64_t(n) < prime &&
std::uint64_t(k) < prime);
maximum = std::max(maximum, n);
}
if (q_ == 0) {
order_ = 0;
return;
}
order_ = prime - 1;
for (auto [factor, exponent] : factorize(order_)) {
(void)exponent;
while (order_ % factor == 0 &&
power(q_, order_ / factor) == 1) {
order_ /= std::uint32_t(factor);
}
}
ordinary_.build(maximum / int(order_), prime);
int residual_limit = std::min<std::uint64_t>(maximum, order_ - 1);
q_factorial_.resize(residual_limit + 1);
inverse_q_factorial_.resize(residual_limit + 1);
q_factorial_[0] = 1;
std::uint32_t q_power = 1;
for (int i = 1; i <= residual_limit; i++) {
q_power = multiply(q_power, q_);
std::uint32_t factor = q_power == 0 ? 1 : modulus_ + 1ULL - q_power;
factor %= modulus_;
assert(factor != 0);
q_factorial_[i] = multiply(q_factorial_[i - 1], factor);
}
inverse_q_factorial_[residual_limit] =
power(q_factorial_[residual_limit], prime - 2);
q_power = power(q_, residual_limit);
std::uint32_t inverse_q = power(q_, prime - 2);
for (int i = residual_limit; i > 0; i--) {
std::uint32_t factor = std::uint32_t(
(std::uint64_t(modulus_) + 1 - q_power) % modulus_);
inverse_q_factorial_[i - 1] =
multiply(inverse_q_factorial_[i], factor);
q_power = multiply(q_power, inverse_q);
}
}
std::uint32_t choose(int n, int k) const {
if (k < 0 || k > n) {
return 0;
}
if (q_ == 0) {
return 1 % modulus_;
}
int high_n = n / int(order_);
int high_k = k / int(order_);
int low_n = n % int(order_);
int low_k = k % int(order_);
if (low_k > low_n) {
return 0;
}
std::uint32_t low =
multiply(q_factorial_[low_n],
multiply(inverse_q_factorial_[low_k],
inverse_q_factorial_[low_n - low_k]));
return multiply(ordinary_.choose(high_n, high_k), low);
}
std::uint32_t order() const { return order_; }
private:
std::uint32_t modulus_ = 1;
std::uint32_t q_ = 0;
std::uint32_t order_ = 0;
prime_binomial_table ordinary_;
std::vector<std::uint32_t> q_factorial_;
std::vector<std::uint32_t> inverse_q_factorial_;
std::uint32_t multiply(std::uint64_t first, std::uint64_t second) const {
return std::uint32_t(first * second % modulus_);
}
std::uint32_t power(std::uint32_t value, std::uint64_t exponent) const {
std::uint32_t result = 1 % modulus_;
while (exponent > 0) {
if (exponent & 1) {
result = multiply(result, value);
}
value = multiply(value, value);
exponent >>= 1;
}
return result;
}
};
} // namespace noya
#endif // NOYA_Q_BINOMIAL_PRIME_HPP
#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <numeric>
#include <utility>
#include <vector>
/// @complexity Time: expected factorization time for p-1 plus
/// O(max(n/order) + min(order,max n) + T) for T queries.
/// Space: O(max(n/order) + min(order,max n)).
/// @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
/// @complexity Time: O(N + log p) preprocessing and O(1) per query.
/// Space: O(N).
namespace noya {
/// @brief Binomial coefficients modulo a runtime prime p for arguments below
/// p. Store factorials and inverse factorials through N; Fermat inversion of
/// N! followed by a backward sweep obtains every inverse using one power.
class prime_binomial_table {
public:
prime_binomial_table() = default;
prime_binomial_table(int maximum, std::uint32_t prime) {
build(maximum, prime);
}
void build(int maximum, std::uint32_t prime) {
assert(maximum >= 0 && std::uint64_t(maximum) < prime);
modulus_ = prime;
factorial_.resize(maximum + 1);
inverse_factorial_.resize(maximum + 1);
factorial_[0] = 1 % prime;
for (int i = 1; i <= maximum; i++) {
factorial_[i] = multiply(factorial_[i - 1], i);
}
inverse_factorial_[maximum] = power(factorial_[maximum], prime - 2);
for (int i = maximum; i > 0; i--) {
inverse_factorial_[i - 1] = multiply(inverse_factorial_[i], i);
}
}
std::uint32_t choose(int n, int k) const {
if (k < 0 || k > n) {
return 0;
}
assert(n < int(factorial_.size()));
return multiply(factorial_[n],
multiply(inverse_factorial_[k],
inverse_factorial_[n - k]));
}
private:
std::uint32_t modulus_ = 1;
std::vector<std::uint32_t> factorial_;
std::vector<std::uint32_t> inverse_factorial_;
std::uint32_t multiply(std::uint64_t first, std::uint64_t second) const {
return std::uint32_t(first * second % modulus_);
}
std::uint32_t power(std::uint32_t value, std::uint64_t exponent) const {
std::uint32_t result = 1 % modulus_;
while (exponent > 0) {
if (exponent & 1) {
result = multiply(result, value);
}
value = multiply(value, value);
exponent >>= 1;
}
return result;
}
};
} // namespace noya
namespace noya {
/// @brief Answer Gaussian binomial coefficients at a fixed q modulo a prime.
/// If d is the multiplicative order of q, q-Lucas gives
/// [n choose k]_q = C(floor(n/d),floor(k/d)) [n mod d choose k mod d]_q.
/// Ordinary factorials handle the first factor, while products of 1-q^i for
/// i<d handle the second. q=0 and q=1 are treated by their direct limits.
class q_binomial_prime_table {
public:
q_binomial_prime_table() = default;
q_binomial_prime_table(
const std::vector<std::pair<int, int>> &queries, std::uint32_t prime,
std::uint32_t q) {
build(queries, prime, q);
}
void build(const std::vector<std::pair<int, int>> &queries,
std::uint32_t prime, std::uint32_t q) {
modulus_ = prime;
q_ = q % prime;
int maximum = 0;
for (auto [n, k] : queries) {
assert(n >= 0 && k >= 0 && std::uint64_t(n) < prime &&
std::uint64_t(k) < prime);
maximum = std::max(maximum, n);
}
if (q_ == 0) {
order_ = 0;
return;
}
order_ = prime - 1;
for (auto [factor, exponent] : factorize(order_)) {
(void)exponent;
while (order_ % factor == 0 &&
power(q_, order_ / factor) == 1) {
order_ /= std::uint32_t(factor);
}
}
ordinary_.build(maximum / int(order_), prime);
int residual_limit = std::min<std::uint64_t>(maximum, order_ - 1);
q_factorial_.resize(residual_limit + 1);
inverse_q_factorial_.resize(residual_limit + 1);
q_factorial_[0] = 1;
std::uint32_t q_power = 1;
for (int i = 1; i <= residual_limit; i++) {
q_power = multiply(q_power, q_);
std::uint32_t factor = q_power == 0 ? 1 : modulus_ + 1ULL - q_power;
factor %= modulus_;
assert(factor != 0);
q_factorial_[i] = multiply(q_factorial_[i - 1], factor);
}
inverse_q_factorial_[residual_limit] =
power(q_factorial_[residual_limit], prime - 2);
q_power = power(q_, residual_limit);
std::uint32_t inverse_q = power(q_, prime - 2);
for (int i = residual_limit; i > 0; i--) {
std::uint32_t factor = std::uint32_t(
(std::uint64_t(modulus_) + 1 - q_power) % modulus_);
inverse_q_factorial_[i - 1] =
multiply(inverse_q_factorial_[i], factor);
q_power = multiply(q_power, inverse_q);
}
}
std::uint32_t choose(int n, int k) const {
if (k < 0 || k > n) {
return 0;
}
if (q_ == 0) {
return 1 % modulus_;
}
int high_n = n / int(order_);
int high_k = k / int(order_);
int low_n = n % int(order_);
int low_k = k % int(order_);
if (low_k > low_n) {
return 0;
}
std::uint32_t low =
multiply(q_factorial_[low_n],
multiply(inverse_q_factorial_[low_k],
inverse_q_factorial_[low_n - low_k]));
return multiply(ordinary_.choose(high_n, high_k), low);
}
std::uint32_t order() const { return order_; }
private:
std::uint32_t modulus_ = 1;
std::uint32_t q_ = 0;
std::uint32_t order_ = 0;
prime_binomial_table ordinary_;
std::vector<std::uint32_t> q_factorial_;
std::vector<std::uint32_t> inverse_q_factorial_;
std::uint32_t multiply(std::uint64_t first, std::uint64_t second) const {
return std::uint32_t(first * second % modulus_);
}
std::uint32_t power(std::uint32_t value, std::uint64_t exponent) const {
std::uint32_t result = 1 % modulus_;
while (exponent > 0) {
if (exponent & 1) {
result = multiply(result, value);
}
value = multiply(value, value);
exponent >>= 1;
}
return result;
}
};
} // namespace noya