232's competitive-programming templates.

PrimePowerCombination (P^E binom)

232-aka-pretty-pi 2024. 8. 8. 21:16

needs Mint (because of static assertion "isPrime")

overflow in constant expression "sz" will throw compile error. (i.e. P = 3, E = 40? power(3, 40) >= power(2, 31))

template <int P, int E>
struct PrimePowerCombination {
  static_assert(P > 0 && isPrime(P) && 0 < E && E < 32);
  constexpr static int sz = [] {
    int sz = 1;
    for (int i = 0; i < E; ++i) {
      sz *= P;
    }
    return sz;
  } ();
  std::array<int, sz> fact;

  constexpr PrimePowerCombination() {
    fact[0] = 1;
    for (int i = 1; i < sz; ++i) {
      fact[i] = (i % P == 0 ? 1 : i) * i64(fact[i - 1]) % sz;
    }
  }

  [[nodiscard]] constexpr std::pair<int, int> fac(int n) const {
    int resr = 1;
    int resp = 0;
    while (n > 0) {
      resr = i64(resr) * fact[n % sz] % sz;
      if (n / sz % 2 == 1 && fact[sz - 1] == sz - 1) {
        if (resr != 0) {
          resr = sz - resr;
        }
      }
      resp += n / P;
      n /= P;
    }
    return {resr, resp};
  }

  [[nodiscard]] constexpr static int inverse(int a) {
    int m = sz;
    int u = 1;
    int v = 0;
    while (m != 0) {
      int t = a / m;
      a -= t * m; std::swap(a, m);
      u -= t * v; std::swap(u, v);
    }
    return u;
  }

  [[nodiscard]] constexpr int C(int n, int r) const {
    if (r < 0 || n < r) {
      return 0;
    }
    auto [r0, p0] = fac(n);
    auto [r1, p1] = fac(r);
    auto [r2, p2] = fac(n - r);
    int p = p0 - p1 - p2;
    if (p >= E) {
      return 0;
    }
    int res = i64(r0) * inverse(r1) % sz * inverse(r2) % sz;
    if (res < 0) {
      res += sz;
    }
    for (int i = 0; i < p; ++i) {
      res = P * res % sz;
    }
    return res;
  }
};