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;
}
};