C++ Library for Competitive Programming
#include "emthrm/math/enumerate_k-th_power.hpp"
$O(N)$
名前 | 戻り値 |
---|---|
std::vector<int> enumerate_kth_power(const int n, const int k, const int m); |
$i^k \bmod m$ ($0 \leq i \leq n$) |
https://yukicoder.me/submissions/623345
#ifndef EMTHRM_MATH_ENUMERATE_K_TH_POWER_HPP_
#define EMTHRM_MATH_ENUMERATE_K_TH_POWER_HPP_
#include <vector>
#include "emthrm/math/mod_pow.hpp"
#include "emthrm/math/prime_sieve.hpp"
namespace emthrm {
std::vector<int> enumerate_kth_power(const int n, const int k, const int m) {
const std::vector<int> smallest_prime_factor = prime_sieve<false>(n);
std::vector<int> res(n + 1, 0);
for (int i = 1; i <= n; ++i) {
if (smallest_prime_factor[i] == i) [[unlikely]] {
res[i] = mod_pow(i, k, m);
} else {
res[i] = static_cast<long long>(res[smallest_prime_factor[i]])
* res[i / smallest_prime_factor[i]] % m;
}
}
return res;
}
} // namespace emthrm
#endif // EMTHRM_MATH_ENUMERATE_K_TH_POWER_HPP_
#line 1 "include/emthrm/math/enumerate_k-th_power.hpp"
#include <vector>
#line 1 "include/emthrm/math/mod_pow.hpp"
namespace emthrm {
long long mod_pow(long long x, long long n, const int m) {
if ((x %= m) < 0) x += m;
long long res = 1;
for (; n > 0; n >>= 1) {
if (n & 1) res = (res * x) % m;
x = (x * x) % m;
}
return res;
}
} // namespace emthrm
#line 1 "include/emthrm/math/prime_sieve.hpp"
#include <numeric>
#line 6 "include/emthrm/math/prime_sieve.hpp"
namespace emthrm {
template <bool GETS_ONLY_PRIME>
std::vector<int> prime_sieve(const int n) {
std::vector<int> smallest_prime_factor(n + 1), prime;
std::iota(smallest_prime_factor.begin(), smallest_prime_factor.end(), 0);
for (int i = 2; i <= n; ++i) {
if (smallest_prime_factor[i] == i) [[unlikely]] prime.emplace_back(i);
for (const int p : prime) {
if (i * p > n || p > smallest_prime_factor[i]) break;
smallest_prime_factor[i * p] = p;
}
}
return GETS_ONLY_PRIME ? prime : smallest_prime_factor;
}
} // namespace emthrm
#line 8 "include/emthrm/math/enumerate_k-th_power.hpp"
namespace emthrm {
std::vector<int> enumerate_kth_power(const int n, const int k, const int m) {
const std::vector<int> smallest_prime_factor = prime_sieve<false>(n);
std::vector<int> res(n + 1, 0);
for (int i = 1; i <= n; ++i) {
if (smallest_prime_factor[i] == i) [[unlikely]] {
res[i] = mod_pow(i, k, m);
} else {
res[i] = static_cast<long long>(res[smallest_prime_factor[i]])
* res[i / smallest_prime_factor[i]] % m;
}
}
return res;
}
} // namespace emthrm