Skip to content

prime_binomial_table.hpp

SECTIONMath INCLUDEnoya/prime_binomial_table.hpp

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.

Verified by binomial_coefficient_prime_mod.

\[ \displaystyle \binom{n}{k} \equiv n!/(k!(n-k)!) \pmod p \]

Implementation

View on GitHub

#ifndef NOYA_PRIME_BINOMIAL_TABLE_HPP
#define NOYA_PRIME_BINOMIAL_TABLE_HPP 1

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

#endif // NOYA_PRIME_BINOMIAL_TABLE_HPP
#include <cassert>
#include <cstdint>
#include <vector>

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