Skip to content

q_binomial_prime.hpp

SECTIONMath INCLUDEnoya/math/q_binomial_prime.hpp

对长 \(n\)、恰有 \(k\) 个 1 的所有 01 串 s,把每个串按 \(q^{\operatorname{inv}(s)}\) 计权再求和,就是 q-组合数;\(\operatorname{inv}(s)\) 是前 1 后 0 的数对个数。例:\(n\)=3、\(k\)=1 时,001、010、100 分别贡献 1、q、\(q^2\)。q=1 时得到普通组合数;q 为素数幂时,也等于 q 元有限域上 \(n\) 维空间中 \(k\) 维子空间的数量。本实现固定 q 和素数模 \(p\),返回代入 q 后模 \(p\) 的数值,不返回各逆序对数对应的方案数;要求 \(0\le n,k<p\le\mathrm{INT\_MAX}\),查询不超过预处理的最大 \(n\)。

\[ \displaystyle \binom{n}{k}_q=\sum_s q^{\operatorname{inv}(s)},\qquad \binom{3}{1}_q=1+q+q^2 \]

Complexity: Time: for \(q > 1\), expected factorization time for \(p - 1\) plus \(O(T + \log^{2} p + N / d + \min (d , N))\) preprocessing; \(O(1)\) per query. \(q = 0\) and \(q = 1\) use \(O(T)\) and \(O(T + N + \log p)\) preprocessing, respectively. Space: for a fresh table, \(O(N / d + \min (d , N) + \log p)\) for \(q > 1\); \(O(1)\) and \(O(N + 1)\) for \(q = 0\) and \(q = 1\), respectively. \(N\) is the largest build argument, \(T\) is the build query count, and \(d\) is the multiplicative order of \(q\). Here \(q\) is reduced modulo \(p\).

AC 记录:q_binomial_coefficient_prime_mod。

跳到代码 · GitHub ↗

Implementation

当前头文件,省略 include guard;依赖见 #include。

/// @complexity Time: for q>1, expected factorization time for p-1 plus
/// O(T + log^2 p + N/d + min(d,N)) preprocessing; O(1) per query.
/// q=0 and q=1 use O(T) and O(T+N+log p) preprocessing, respectively.
/// Space: for a fresh table, O(N/d + min(d,N) + log p) for q>1;
/// O(1) and O(N+1) for q=0 and q=1, respectively.
/// N is the largest build argument, T is the build query count, and d is
/// the multiplicative order of q. Here q is reduced modulo p.

#include "noya/math/factorize.hpp"
#include "noya/math/prime_binomial_table.hpp"

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

namespace noya {

/// @brief Sum q^inv(s) over length-n binary strings with k ones, modulo p.
/// inv(s) counts pairs i<j with s[i]=1 and s[j]=0. This is the Gaussian
/// binomial [n choose k]_q; e.g. [3 choose 1]_q = 1+q+q^2.
/// For prime-power q it also counts k-dimensional subspaces of F_q^n.
/// Returns evaluation at one fixed q, not the polynomial's coefficients.
/// Require prime 2<=p<=INT_MAX and 0<=n,k<p in the build query list.
/// Later queries must have 0<=n<=max build n; choose returns 0 for k<0 or k>n.
/// 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>> &que,
                         std::uint32_t p, std::uint32_t q) {
    build(que, p, q);
  }

  void build(const std::vector<std::pair<int, int>> &que, std::uint32_t p,
             std::uint32_t q) {
    assert(2 <= p && p <= INT_MAX);
    md = p;
    q_ = q % p;
    int mx = 0;
    for (auto [n, k] : que) {
      assert(n >= 0 && k >= 0 && std::uint64_t(n) < p && std::uint64_t(k) < p);
      mx = std::max(mx, n);
    }
    if (q_ == 0) {
      od_ = 0;
      return;
    }
    if (q_ == 1) {
      od_ = 1;
      ord.build(mx, p);
      return;
    }
    od_ = p - 1;
    for (auto [fct, exp] : factorize(od_)) {
      (void)exp;
      while (od_ % fct == 0 && power(q_, od_ / fct) == 1) {
        od_ /= std::uint32_t(fct);
      }
    }

    ord.build(mx / int(od_), p);
    int lim = std::min<std::uint64_t>(mx, od_ - 1);
    fac.resize(lim + 1);
    ifc.resize(lim + 1);
    fac[0] = 1;
    std::uint32_t qp = 1;
    for (int i = 1; i <= lim; i++) {
      qp = multiply(qp, q_);
      // i < od_, so q^i is neither zero nor one.
      std::uint32_t fct = md + 1 - qp;
      fac[i] = multiply(fac[i - 1], fct);
    }
    ifc[lim] = power(fac[lim], p - 2);
    std::uint32_t iq = power(q_, od_ - 1);
    for (int i = lim; i > 0; i--) {
      std::uint32_t fct = md + 1 - qp;
      ifc[i - 1] = multiply(ifc[i], fct);
      qp = multiply(qp, iq);
    }
  }

  std::uint32_t choose(int n, int k) const {
    if (k < 0 || k > n) {
      return 0;
    }
    if (q_ == 0) {
      return 1 % md;
    }
    if (q_ == 1)
      return ord.choose(n, k);
    int hn = n / int(od_);
    int hk = k / int(od_);
    int ln = n % int(od_);
    int lk = k % int(od_);
    if (lk > ln) {
      return 0;
    }
    std::uint32_t low = multiply(fac[ln], multiply(ifc[lk], ifc[ln - lk]));
    return multiply(ord.choose(hn, hk), low);
  }

  std::uint32_t order() const { return od_; }

private:
  std::uint32_t md = 1;
  std::uint32_t q_ = 0;
  std::uint32_t od_ = 0;
  prime_binomial_table ord;
  std::vector<std::uint32_t> fac;
  std::vector<std::uint32_t> ifc;

  std::uint32_t multiply(std::uint64_t a, std::uint64_t b) const {
    return std::uint32_t(a * b % md);
  }

  std::uint32_t power(std::uint32_t val, std::uint64_t exp) const {
    std::uint32_t res = 1 % md;
    while (exp > 0) {
      if (exp & 1) {
        res = multiply(res, val);
      }
      val = multiply(val, val);
      exp >>= 1;
    }
    return res;
  }
};

} // namespace noya

 // NOYA_Q_BINOMIAL_PRIME_HPP
#ifndef NOYA_Q_BINOMIAL_PRIME_HPP
#define NOYA_Q_BINOMIAL_PRIME_HPP 1

/// @complexity Time: for q>1, expected factorization time for p-1 plus
/// O(T + log^2 p + N/d + min(d,N)) preprocessing; O(1) per query.
/// q=0 and q=1 use O(T) and O(T+N+log p) preprocessing, respectively.
/// Space: for a fresh table, O(N/d + min(d,N) + log p) for q>1;
/// O(1) and O(N+1) for q=0 and q=1, respectively.
/// N is the largest build argument, T is the build query count, and d is
/// the multiplicative order of q. Here q is reduced modulo p.

#include "noya/math/factorize.hpp"
#include "noya/math/prime_binomial_table.hpp"

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

namespace noya {

/// @brief Sum q^inv(s) over length-n binary strings with k ones, modulo p.
/// inv(s) counts pairs i<j with s[i]=1 and s[j]=0. This is the Gaussian
/// binomial [n choose k]_q; e.g. [3 choose 1]_q = 1+q+q^2.
/// For prime-power q it also counts k-dimensional subspaces of F_q^n.
/// Returns evaluation at one fixed q, not the polynomial's coefficients.
/// Require prime 2<=p<=INT_MAX and 0<=n,k<p in the build query list.
/// Later queries must have 0<=n<=max build n; choose returns 0 for k<0 or k>n.
/// 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>> &que,
                         std::uint32_t p, std::uint32_t q) {
    build(que, p, q);
  }

  void build(const std::vector<std::pair<int, int>> &que, std::uint32_t p,
             std::uint32_t q) {
    assert(2 <= p && p <= INT_MAX);
    md = p;
    q_ = q % p;
    int mx = 0;
    for (auto [n, k] : que) {
      assert(n >= 0 && k >= 0 && std::uint64_t(n) < p && std::uint64_t(k) < p);
      mx = std::max(mx, n);
    }
    if (q_ == 0) {
      od_ = 0;
      return;
    }
    if (q_ == 1) {
      od_ = 1;
      ord.build(mx, p);
      return;
    }
    od_ = p - 1;
    for (auto [fct, exp] : factorize(od_)) {
      (void)exp;
      while (od_ % fct == 0 && power(q_, od_ / fct) == 1) {
        od_ /= std::uint32_t(fct);
      }
    }

    ord.build(mx / int(od_), p);
    int lim = std::min<std::uint64_t>(mx, od_ - 1);
    fac.resize(lim + 1);
    ifc.resize(lim + 1);
    fac[0] = 1;
    std::uint32_t qp = 1;
    for (int i = 1; i <= lim; i++) {
      qp = multiply(qp, q_);
      // i < od_, so q^i is neither zero nor one.
      std::uint32_t fct = md + 1 - qp;
      fac[i] = multiply(fac[i - 1], fct);
    }
    ifc[lim] = power(fac[lim], p - 2);
    std::uint32_t iq = power(q_, od_ - 1);
    for (int i = lim; i > 0; i--) {
      std::uint32_t fct = md + 1 - qp;
      ifc[i - 1] = multiply(ifc[i], fct);
      qp = multiply(qp, iq);
    }
  }

  std::uint32_t choose(int n, int k) const {
    if (k < 0 || k > n) {
      return 0;
    }
    if (q_ == 0) {
      return 1 % md;
    }
    if (q_ == 1)
      return ord.choose(n, k);
    int hn = n / int(od_);
    int hk = k / int(od_);
    int ln = n % int(od_);
    int lk = k % int(od_);
    if (lk > ln) {
      return 0;
    }
    std::uint32_t low = multiply(fac[ln], multiply(ifc[lk], ifc[ln - lk]));
    return multiply(ord.choose(hn, hk), low);
  }

  std::uint32_t order() const { return od_; }

private:
  std::uint32_t md = 1;
  std::uint32_t q_ = 0;
  std::uint32_t od_ = 0;
  prime_binomial_table ord;
  std::vector<std::uint32_t> fac;
  std::vector<std::uint32_t> ifc;

  std::uint32_t multiply(std::uint64_t a, std::uint64_t b) const {
    return std::uint32_t(a * b % md);
  }

  std::uint32_t power(std::uint32_t val, std::uint64_t exp) const {
    std::uint32_t res = 1 % md;
    while (exp > 0) {
      if (exp & 1) {
        res = multiply(res, val);
      }
      val = multiply(val, val);
      exp >>= 1;
    }
    return res;
  }
};

} // namespace noya

#endif // NOYA_Q_BINOMIAL_PRIME_HPP
/// @complexity Time: for q>1, expected factorization time for p-1 plus
/// O(T + log^2 p + N/d + min(d,N)) preprocessing; O(1) per query.
/// q=0 and q=1 use O(T) and O(T+N+log p) preprocessing, respectively.
/// Space: for a fresh table, O(N/d + min(d,N) + log p) for q>1;
/// O(1) and O(N+1) for q=0 and q=1, respectively.
/// N is the largest build argument, T is the build query count, and d is
/// the multiplicative order of q. Here q is reduced modulo p.

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

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

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) {
      val = multiply_mod(val, val, n);
      // Add c modulo n without overflowing 64 bits.
      return val >= n - c ? val - (n - c) : val + c;
    };
    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

/// @complexity Time: O(N + log p) preprocessing and O(1) per query.
/// Space: O(N).

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

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 mx, std::uint32_t p) { build(mx, p); }

  void build(int mx, std::uint32_t p) {
    assert(mx >= 0 && std::uint64_t(mx) < p);
    md = p;
    fac.resize(mx + 1);
    ifc.resize(mx + 1);
    fac[0] = 1 % p;
    for (int i = 1; i <= mx; i++) {
      fac[i] = multiply(fac[i - 1], i);
    }
    ifc[mx] = power(fac[mx], p - 2);
    for (int i = mx; i > 0; i--) {
      ifc[i - 1] = multiply(ifc[i], i);
    }
  }

  std::uint32_t choose(int n, int k) const {
    if (k < 0 || k > n) {
      return 0;
    }
    assert(n < int(fac.size()));
    return multiply(fac[n], multiply(ifc[k], ifc[n - k]));
  }

private:
  std::uint32_t md = 1;
  std::vector<std::uint32_t> fac;
  std::vector<std::uint32_t> ifc;

  std::uint32_t multiply(std::uint64_t a, std::uint64_t b) const {
    return std::uint32_t(a * b % md);
  }

  std::uint32_t power(std::uint32_t val, std::uint64_t exp) const {
    std::uint32_t res = 1 % md;
    while (exp > 0) {
      if (exp & 1) {
        res = multiply(res, val);
      }
      val = multiply(val, val);
      exp >>= 1;
    }
    return res;
  }
};

} // namespace noya

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

namespace noya {

/// @brief Sum q^inv(s) over length-n binary strings with k ones, modulo p.
/// inv(s) counts pairs i<j with s[i]=1 and s[j]=0. This is the Gaussian
/// binomial [n choose k]_q; e.g. [3 choose 1]_q = 1+q+q^2.
/// For prime-power q it also counts k-dimensional subspaces of F_q^n.
/// Returns evaluation at one fixed q, not the polynomial's coefficients.
/// Require prime 2<=p<=INT_MAX and 0<=n,k<p in the build query list.
/// Later queries must have 0<=n<=max build n; choose returns 0 for k<0 or k>n.
/// 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>> &que,
                         std::uint32_t p, std::uint32_t q) {
    build(que, p, q);
  }

  void build(const std::vector<std::pair<int, int>> &que, std::uint32_t p,
             std::uint32_t q) {
    assert(2 <= p && p <= INT_MAX);
    md = p;
    q_ = q % p;
    int mx = 0;
    for (auto [n, k] : que) {
      assert(n >= 0 && k >= 0 && std::uint64_t(n) < p && std::uint64_t(k) < p);
      mx = std::max(mx, n);
    }
    if (q_ == 0) {
      od_ = 0;
      return;
    }
    if (q_ == 1) {
      od_ = 1;
      ord.build(mx, p);
      return;
    }
    od_ = p - 1;
    for (auto [fct, exp] : factorize(od_)) {
      (void)exp;
      while (od_ % fct == 0 && power(q_, od_ / fct) == 1) {
        od_ /= std::uint32_t(fct);
      }
    }

    ord.build(mx / int(od_), p);
    int lim = std::min<std::uint64_t>(mx, od_ - 1);
    fac.resize(lim + 1);
    ifc.resize(lim + 1);
    fac[0] = 1;
    std::uint32_t qp = 1;
    for (int i = 1; i <= lim; i++) {
      qp = multiply(qp, q_);
      // i < od_, so q^i is neither zero nor one.
      std::uint32_t fct = md + 1 - qp;
      fac[i] = multiply(fac[i - 1], fct);
    }
    ifc[lim] = power(fac[lim], p - 2);
    std::uint32_t iq = power(q_, od_ - 1);
    for (int i = lim; i > 0; i--) {
      std::uint32_t fct = md + 1 - qp;
      ifc[i - 1] = multiply(ifc[i], fct);
      qp = multiply(qp, iq);
    }
  }

  std::uint32_t choose(int n, int k) const {
    if (k < 0 || k > n) {
      return 0;
    }
    if (q_ == 0) {
      return 1 % md;
    }
    if (q_ == 1)
      return ord.choose(n, k);
    int hn = n / int(od_);
    int hk = k / int(od_);
    int ln = n % int(od_);
    int lk = k % int(od_);
    if (lk > ln) {
      return 0;
    }
    std::uint32_t low = multiply(fac[ln], multiply(ifc[lk], ifc[ln - lk]));
    return multiply(ord.choose(hn, hk), low);
  }

  std::uint32_t order() const { return od_; }

private:
  std::uint32_t md = 1;
  std::uint32_t q_ = 0;
  std::uint32_t od_ = 0;
  prime_binomial_table ord;
  std::vector<std::uint32_t> fac;
  std::vector<std::uint32_t> ifc;

  std::uint32_t multiply(std::uint64_t a, std::uint64_t b) const {
    return std::uint32_t(a * b % md);
  }

  std::uint32_t power(std::uint32_t val, std::uint64_t exp) const {
    std::uint32_t res = 1 % md;
    while (exp > 0) {
      if (exp & 1) {
        res = multiply(res, val);
      }
      val = multiply(val, val);
      exp >>= 1;
    }
    return res;
  }
};

} // namespace noya

 // NOYA_Q_BINOMIAL_PRIME_HPP