Skip to content

stirling_small_prime.hpp

SECTIONMath INCLUDEnoya/stirling_small_prime.hpp

Query signed first-kind and second-kind Stirling numbers modulo a small fixed prime for indices as large as 10^18. Over F_p, the identities (x)_p=x^p-x and x^p=x split a large falling factorial or power into p-sized blocks. Repeating the split leaves one small Stirling number with indices below p and one binomial coefficient. Lucas' theorem evaluates the latter digit by digit. The constructor can omit either Stirling table when only one kind is needed.

Verified by stirling_number_of_the_first_kind_small_p_large_n, stirling_number_of_the_second_kind_small_p_large_n.

\[ \displaystyle x^n=\sum_{k=0}^n S(n,k)x^{\underline{k}} \]

Implementation

View on GitHub

#ifndef NOYA_STIRLING_SMALL_PRIME_HPP
#define NOYA_STIRLING_SMALL_PRIME_HPP 1

/// @complexity Time: O(p^2) preprocessing and O(log_p n) per query.
/// Space: O(p^2) for each requested Stirling kind.

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

namespace noya {

/// @brief Query signed first-kind and second-kind Stirling numbers modulo a
/// small fixed prime for indices as large as 10^18. Over F_p, the identities
/// (x)_p=x^p-x and x^p=x split a large falling factorial or power into p-sized
/// blocks. Repeating the split leaves one small Stirling number with indices
/// below p and one binomial coefficient. Lucas' theorem evaluates the latter
/// digit by digit. The constructor can omit either Stirling table when only one
/// kind is needed.
class stirling_small_prime_table {
public:
  explicit stirling_small_prime_table(std::uint32_t prime,
                                      bool build_first_kind = true,
                                      bool build_second_kind = true)
      : prime_(prime) {
    assert(prime >= 2);
    build_binomial();
    if (build_first_kind) {
      build_first();
    }
    if (build_second_kind) {
      build_second();
    }
  }

  std::uint32_t first_kind(std::uint64_t n, std::uint64_t k) const {
    assert(!first_.empty());
    if (k > n) {
      return 0;
    }
    std::uint64_t blocks = n / prime_;
    int remainder = int(n % prime_);
    if (blocks > k) {
      return 0;
    }
    std::uint64_t quotient = (k - blocks) / (prime_ - 1);
    int small_k = int((k - blocks) % (prime_ - 1));
    if (small_k == 0 && remainder > 0) {
      if (quotient == 0) {
        return 0;
      }
      small_k = int(prime_ - 1);
      quotient--;
    }
    if (blocks < quotient || small_k > remainder) {
      return 0;
    }
    std::uint32_t result = multiply(
        lucas_choose(blocks, quotient), first_[remainder][small_k]);
    if (((blocks + quotient) & 1) != 0 && result != 0) {
      result = prime_ - result;
    }
    return result;
  }

  std::uint32_t second_kind(std::uint64_t n, std::uint64_t k) const {
    assert(!second_.empty());
    if (k > n) {
      return 0;
    }
    if (n == 0) {
      return k == 0 ? 1 : 0;
    }
    std::uint64_t blocks = k / prime_;
    int small_k = int(k % prime_);
    if (n < blocks) {
      return 0;
    }
    std::uint64_t quotient = (n - blocks) / (prime_ - 1);
    int remainder = int((n - blocks) % (prime_ - 1));
    if (remainder == 0) {
      if (quotient == 0) {
        return 0;
      }
      remainder = int(prime_ - 1);
      quotient--;
    }
    if (small_k > remainder) {
      return 0;
    }
    if (remainder < int(prime_ - 1)) {
      return multiply(lucas_choose(quotient, blocks),
                      second_[remainder][small_k]);
    }
    if (small_k == 0) {
      return blocks == 0 ? 0 : lucas_choose(quotient, blocks - 1);
    }
    return multiply(lucas_choose(quotient, blocks),
                    second_[prime_ - 1][small_k]);
  }

private:
  std::uint32_t prime_;
  std::vector<std::vector<std::uint32_t>> binomial_;
  std::vector<std::vector<std::uint32_t>> first_;
  std::vector<std::vector<std::uint32_t>> second_;

  std::uint32_t multiply(std::uint64_t first,
                         std::uint64_t second) const {
    return std::uint32_t(first * second % prime_);
  }

  std::uint32_t lucas_choose(std::uint64_t n, std::uint64_t k) const {
    if (k > n) {
      return 0;
    }
    std::uint32_t result = 1;
    while (n > 0 || k > 0) {
      int upper = int(n % prime_);
      int lower = int(k % prime_);
      if (lower > upper) {
        return 0;
      }
      result = multiply(result, binomial_[upper][lower]);
      n /= prime_;
      k /= prime_;
    }
    return result;
  }

  void build_binomial() {
    binomial_.resize(prime_);
    binomial_[0] = {1};
    for (int n = 1; n < int(prime_); n++) {
      binomial_[n] = binomial_[n - 1];
      binomial_[n].push_back(0);
      for (int k = 1; k <= n; k++) {
        binomial_[n][k] += binomial_[n - 1][k - 1];
        if (binomial_[n][k] >= prime_) {
          binomial_[n][k] -= prime_;
        }
      }
    }
  }

  void build_first() {
    first_.resize(prime_);
    first_[0] = {1};
    for (int n = 1; n < int(prime_); n++) {
      first_[n].assign(n + 1, 0);
      for (int k = 0; k <= n; k++) {
        if (k > 0) {
          first_[n][k] += first_[n - 1][k - 1];
        }
        if (k < n) {
          first_[n][k] +=
              multiply(first_[n - 1][k], prime_ - n + 1);
        }
        if (first_[n][k] >= prime_) {
          first_[n][k] -= prime_;
        }
      }
    }
  }

  void build_second() {
    second_.resize(prime_);
    second_[0] = {1};
    for (int n = 1; n < int(prime_); n++) {
      second_[n].assign(n + 1, 0);
      for (int k = 0; k <= n; k++) {
        if (k > 0) {
          second_[n][k] += second_[n - 1][k - 1];
        }
        if (k < n) {
          second_[n][k] += multiply(second_[n - 1][k], k);
        }
        if (second_[n][k] >= prime_) {
          second_[n][k] -= prime_;
        }
      }
    }
  }
};

} // namespace noya

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

/// @complexity Time: O(p^2) preprocessing and O(log_p n) per query.
/// Space: O(p^2) for each requested Stirling kind.

namespace noya {

/// @brief Query signed first-kind and second-kind Stirling numbers modulo a
/// small fixed prime for indices as large as 10^18. Over F_p, the identities
/// (x)_p=x^p-x and x^p=x split a large falling factorial or power into p-sized
/// blocks. Repeating the split leaves one small Stirling number with indices
/// below p and one binomial coefficient. Lucas' theorem evaluates the latter
/// digit by digit. The constructor can omit either Stirling table when only one
/// kind is needed.
class stirling_small_prime_table {
public:
  explicit stirling_small_prime_table(std::uint32_t prime,
                                      bool build_first_kind = true,
                                      bool build_second_kind = true)
      : prime_(prime) {
    assert(prime >= 2);
    build_binomial();
    if (build_first_kind) {
      build_first();
    }
    if (build_second_kind) {
      build_second();
    }
  }

  std::uint32_t first_kind(std::uint64_t n, std::uint64_t k) const {
    assert(!first_.empty());
    if (k > n) {
      return 0;
    }
    std::uint64_t blocks = n / prime_;
    int remainder = int(n % prime_);
    if (blocks > k) {
      return 0;
    }
    std::uint64_t quotient = (k - blocks) / (prime_ - 1);
    int small_k = int((k - blocks) % (prime_ - 1));
    if (small_k == 0 && remainder > 0) {
      if (quotient == 0) {
        return 0;
      }
      small_k = int(prime_ - 1);
      quotient--;
    }
    if (blocks < quotient || small_k > remainder) {
      return 0;
    }
    std::uint32_t result = multiply(
        lucas_choose(blocks, quotient), first_[remainder][small_k]);
    if (((blocks + quotient) & 1) != 0 && result != 0) {
      result = prime_ - result;
    }
    return result;
  }

  std::uint32_t second_kind(std::uint64_t n, std::uint64_t k) const {
    assert(!second_.empty());
    if (k > n) {
      return 0;
    }
    if (n == 0) {
      return k == 0 ? 1 : 0;
    }
    std::uint64_t blocks = k / prime_;
    int small_k = int(k % prime_);
    if (n < blocks) {
      return 0;
    }
    std::uint64_t quotient = (n - blocks) / (prime_ - 1);
    int remainder = int((n - blocks) % (prime_ - 1));
    if (remainder == 0) {
      if (quotient == 0) {
        return 0;
      }
      remainder = int(prime_ - 1);
      quotient--;
    }
    if (small_k > remainder) {
      return 0;
    }
    if (remainder < int(prime_ - 1)) {
      return multiply(lucas_choose(quotient, blocks),
                      second_[remainder][small_k]);
    }
    if (small_k == 0) {
      return blocks == 0 ? 0 : lucas_choose(quotient, blocks - 1);
    }
    return multiply(lucas_choose(quotient, blocks),
                    second_[prime_ - 1][small_k]);
  }

private:
  std::uint32_t prime_;
  std::vector<std::vector<std::uint32_t>> binomial_;
  std::vector<std::vector<std::uint32_t>> first_;
  std::vector<std::vector<std::uint32_t>> second_;

  std::uint32_t multiply(std::uint64_t first,
                         std::uint64_t second) const {
    return std::uint32_t(first * second % prime_);
  }

  std::uint32_t lucas_choose(std::uint64_t n, std::uint64_t k) const {
    if (k > n) {
      return 0;
    }
    std::uint32_t result = 1;
    while (n > 0 || k > 0) {
      int upper = int(n % prime_);
      int lower = int(k % prime_);
      if (lower > upper) {
        return 0;
      }
      result = multiply(result, binomial_[upper][lower]);
      n /= prime_;
      k /= prime_;
    }
    return result;
  }

  void build_binomial() {
    binomial_.resize(prime_);
    binomial_[0] = {1};
    for (int n = 1; n < int(prime_); n++) {
      binomial_[n] = binomial_[n - 1];
      binomial_[n].push_back(0);
      for (int k = 1; k <= n; k++) {
        binomial_[n][k] += binomial_[n - 1][k - 1];
        if (binomial_[n][k] >= prime_) {
          binomial_[n][k] -= prime_;
        }
      }
    }
  }

  void build_first() {
    first_.resize(prime_);
    first_[0] = {1};
    for (int n = 1; n < int(prime_); n++) {
      first_[n].assign(n + 1, 0);
      for (int k = 0; k <= n; k++) {
        if (k > 0) {
          first_[n][k] += first_[n - 1][k - 1];
        }
        if (k < n) {
          first_[n][k] +=
              multiply(first_[n - 1][k], prime_ - n + 1);
        }
        if (first_[n][k] >= prime_) {
          first_[n][k] -= prime_;
        }
      }
    }
  }

  void build_second() {
    second_.resize(prime_);
    second_[0] = {1};
    for (int n = 1; n < int(prime_); n++) {
      second_[n].assign(n + 1, 0);
      for (int k = 0; k <= n; k++) {
        if (k > 0) {
          second_[n][k] += second_[n - 1][k - 1];
        }
        if (k < n) {
          second_[n][k] += multiply(second_[n - 1][k], k);
        }
        if (second_[n][k] >= prime_) {
          second_[n][k] -= prime_;
        }
      }
    }
  }
};

} // namespace noya