cp-library
競技プログラミング用 C++ ライブラリ
剰余を取った冪の列
components/modular_power_sequence.hpp · GitHub で見る
共通ページ
乗法畳み込みで使用する内部部品です。通常は 乗法畳み込み を使ってください。
modular_power_sequence_detail::generate(n, root, mod) は、root^i mod mod を i = 0, ..., n-1 の順に並べた std::vector<std::uint32_t> を返します。n は std::size_t、root・mod は std::uint32_t で、mod >= 1 とします。空の列にも対応し、法 1 では全要素が 0 です。入力の法は素数でなくても構いません。
x86 の AVX2 が使え、法が \(2^{31}\) 未満なら 8 要素ずつ生成します。それ以外の環境・法では通常の整数演算を使います。
時間・返り値を含む空間は \(O(n)\) です。空列は定数時間・空間です。
依存ライブラリ
コード
#ifndef MODULAR_POWER_SEQUENCE_HPP
#define MODULAR_POWER_SEQUENCE_HPP
#include "barrett32.hpp"
#include <vector>
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
#include <immintrin.h>
#endif
namespace modular_power_sequence_detail {
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
__attribute__((target("avx2")))
inline void generate(std::uint32_t *out, std::size_t n, std::uint32_t root, const Barrett32 &arithmetic) {
alignas(32) std::uint32_t first[8];
first[0] = 1 % arithmetic.mod;
for (int j = 1; j < 8; ++j) first[j] = arithmetic.mul(first[j - 1], root);
const std::uint32_t step = arithmetic.mul(first[7], root);
const auto factor = _mm256_set1_epi32(step), quotient = _mm256_set1_epi32((std::uint64_t(step) << 32) / arithmetic.mod);
const auto mod = _mm256_set1_epi32(arithmetic.mod);
auto x = _mm256_load_si256(reinterpret_cast<const __m256i *>(first));
std::size_t i = 0;
for (; i + 8 <= n; i += 8) {
_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i), x);
const auto even = _mm256_srli_epi64(_mm256_mul_epu32(x, quotient), 32);
const auto odd = _mm256_mul_epu32(_mm256_srli_epi64(x, 32), quotient);
const auto q = _mm256_blend_epi32(even, odd, 0xaa);
x = _mm256_sub_epi32(_mm256_mullo_epi32(x, factor), _mm256_mullo_epi32(q, mod));
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, mod));
}
std::uint32_t value = _mm256_extract_epi32(x, 0);
for (; i < n; ++i) { out[i] = value; value = arithmetic.mul(value, root); }
}
#endif
inline std::vector<std::uint32_t> generate(std::size_t n, std::uint32_t root, std::uint32_t mod) {
Barrett32 arithmetic(mod);
std::vector<std::uint32_t> result(n);
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
if (n >= 256 && mod < (std::uint32_t(1) << 31) && __builtin_cpu_supports("avx2")) {
generate(result.data(), n, root, arithmetic);
return result;
}
#endif
std::uint32_t value = 1 % mod;
for (std::size_t i = 0; i < n; ++i) { result[i] = value; value = arithmetic.mul(value, root); }
return result;
}
}
#endif
#line 1 "components/modular_power_sequence.hpp"
#line 1 "components/barrett32.hpp"
#include <cassert>
#include <cstdint>
struct Barrett32 {
std::uint32_t mod;
std::uint64_t reciprocal;
explicit Barrett32(std::uint32_t m) : mod(m), reciprocal(m > 1 ? std::uint64_t((__uint128_t(1) << 64) / m) : 0) {
assert(m > 0);
}
std::uint32_t reduce(std::uint64_t x) const {
if (mod == 1) return 0;
std::uint64_t q = (__uint128_t(x) * reciprocal) >> 64, r = x - q * mod;
return r < mod ? r : r - mod;
}
std::uint32_t mul(std::uint32_t a, std::uint32_t b) const { return reduce(std::uint64_t(a) * b); }
std::uint32_t pow(std::uint32_t a, std::uint64_t e) const {
std::uint32_t r = mod == 1 ? 0 : 1;
for (; e; e >>= 1, a = mul(a, a)) if (e & 1) r = mul(r, a);
return r;
}
};
#line 5 "components/modular_power_sequence.hpp"
#include <vector>
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
#include <immintrin.h>
#endif
namespace modular_power_sequence_detail {
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
__attribute__((target("avx2")))
inline void generate(std::uint32_t *out, std::size_t n, std::uint32_t root, const Barrett32 &arithmetic) {
alignas(32) std::uint32_t first[8];
first[0] = 1 % arithmetic.mod;
for (int j = 1; j < 8; ++j) first[j] = arithmetic.mul(first[j - 1], root);
const std::uint32_t step = arithmetic.mul(first[7], root);
const auto factor = _mm256_set1_epi32(step), quotient = _mm256_set1_epi32((std::uint64_t(step) << 32) / arithmetic.mod);
const auto mod = _mm256_set1_epi32(arithmetic.mod);
auto x = _mm256_load_si256(reinterpret_cast<const __m256i *>(first));
std::size_t i = 0;
for (; i + 8 <= n; i += 8) {
_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i), x);
const auto even = _mm256_srli_epi64(_mm256_mul_epu32(x, quotient), 32);
const auto odd = _mm256_mul_epu32(_mm256_srli_epi64(x, 32), quotient);
const auto q = _mm256_blend_epi32(even, odd, 0xaa);
x = _mm256_sub_epi32(_mm256_mullo_epi32(x, factor), _mm256_mullo_epi32(q, mod));
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, mod));
}
std::uint32_t value = _mm256_extract_epi32(x, 0);
for (; i < n; ++i) { out[i] = value; value = arithmetic.mul(value, root); }
}
#endif
inline std::vector<std::uint32_t> generate(std::size_t n, std::uint32_t root, std::uint32_t mod) {
Barrett32 arithmetic(mod);
std::vector<std::uint32_t> result(n);
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
if (n >= 256 && mod < (std::uint32_t(1) << 31) && __builtin_cpu_supports("avx2")) {
generate(result.data(), n, root, arithmetic);
return result;
}
#endif
std::uint32_t value = 1 % mod;
for (std::size_t i = 0; i < n; ++i) { result[i] = value; value = arithmetic.mul(value, root); }
return result;
}
}