cp-library
競技プログラミング用 C++ ライブラリ
GCD-plus・LCM-plus 畳み込み
convolution/gcd_lcm_convolution.hpp · GitHub で見る
共通ページ
#include "lib/convolution/gcd_lcm_convolution.hpp"
std::vector<long long> a{0, 1, 2, 3}, b{0, 4, 5, 6};
auto c = gcd_convolution(a, b); // {0,62,10,18}
auto d = lcm_convolution(a, b); // {0,4,23,36}
| 操作 | 内容・返り値 | 計算量 |
|---|---|---|
gcd_convolution(a,b) |
\(\displaystyle c_k=\sum_{\gcd(i,j)=k}a_ib_j\) | \(O(N\log\log N)\) |
lcm_convolution(a,b) |
\(\displaystyle c_k=\sum_{\operatorname{lcm}(i,j)=k}a_ib_j\) | \(O(N\log\log N)\) |
値 1 からの配列を、先頭にダミーの 0 を置いて渡します。例えば a[2] は添字 2 の重みです。GCD 版では積 a[2]*b[3] を c[1] に加え、LCM 版では添字 6 が出力範囲内にある場合だけ加えます。入力を変更せず、入力と同じ係数型の配列を返します。
gcd_convolution は \(\displaystyle c_k = \sum_{\gcd(i,j)=k} a_i b_j\)、lcm_convolution は \(\displaystyle c_k = \sum_{\operatorname{lcm}(i,j)=k} a_i b_j\) を計算します。
係数は環の要素とし、\(a_0 = b_0 = 0\) で使います。出力長 \(N = \max(\lvert a\rvert,\lvert b\rvert)\) で、LCM が \(N\) 以上になる項は出力しません。
時間 \(O(N \log \log N)\)、追加空間 \(O(N)\) です。
依存ライブラリ
コード
#ifndef GCD_LCM_CONVOLUTION_HPP
#define GCD_LCM_CONVOLUTION_HPP
#include "../components/gcd_convolution_mint.hpp"
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
#include <vector>
namespace gcd_lcm_convolution_detail {
template <bool gcd, class T>
std::vector<T> multiply(std::vector<T> a, std::vector<T> b) {
if (a.empty() || b.empty()) return {};
assert(a[0] == T(0) && b[0] == T(0));
std::size_t n = std::max(a.size(), b.size());
a.resize(n, T(0));
b.resize(n, T(0));
auto ps = prime_sieve(n);
if constexpr (gcd && std::is_same_v<T, mint_detail::Mint>) {
if (n >= 4096 && n - 1 <= UINT32_MAX) {
gcd_convolution_mint_detail::multiply(a, b, ps);
return a;
}
}
divisor_zeta_detail::for_each<false, gcd>(n, ps, ps.size(), [&](std::size_t to, std::size_t from) {
a[to] += a[from], b[to] += b[from];
});
for (std::size_t i = 1; i < n; ++i) a[i] *= b[i];
divisor_zeta<true, gcd>(a, ps);
return a;
}
}
template <class T>
std::vector<T> gcd_convolution(std::vector<T> a, std::vector<T> b) {
return gcd_lcm_convolution_detail::multiply<true>(std::move(a), std::move(b));
}
template <class T>
std::vector<T> lcm_convolution(std::vector<T> a, std::vector<T> b) {
return gcd_lcm_convolution_detail::multiply<false>(std::move(a), std::move(b));
}
#endif
#line 1 "convolution/gcd_lcm_convolution.hpp"
#line 1 "components/gcd_convolution_mint.hpp"
#line 1 "components/divisor_zeta.hpp"
#line 1 "math/prime_sieve.hpp"
#include <algorithm>
#include <array>
#include <bit>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <vector>
inline std::vector<std::size_t> prime_sieve(std::size_t n);
template<class F> void for_each_prime(std::size_t n, F emit) {
if (n <= 2) return;
emit(std::size_t(2));
if (n <= 65537) {
const std::size_t size = n / 2;
std::vector<std::uint8_t> prime(size, 1);
prime[0] = 0;
for (std::size_t p = 3; p <= (n - 1) / p; p += 2) {
if (!prime[p / 2]) continue;
for (std::size_t j = p * p / 2; j < size; j += p) prime[j] = 0;
}
std::size_t i = 0;
#if defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
for (; i + 8 <= size; i += 8) {
std::uint64_t word;
std::memcpy(&word, prime.data() + i, sizeof(word));
while (word) {
emit(2 * (i + (std::countr_zero(word) >> 3)) + 1);
word &= word - 1;
}
}
#endif
for (; i < size; ++i) if (prime[i]) emit(2 * i + 1);
return;
}
emit(3), emit(5);
constexpr std::array<unsigned, 8> residues{1, 7, 11, 13, 17, 19, 23, 29};
constexpr auto index = [&] {
std::array<unsigned, 30> index{};
for (unsigned i = 0; i < 8; ++i) index[residues[i]] = i;
return index;
}();
constexpr std::size_t period = 7 * 11 * 13, block_size = period * 32;
static const auto pattern = [&] {
std::array<std::uint8_t, period> pattern{};
for (std::size_t i = 0; i < period; ++i)
for (unsigned d = 0; d < 8; ++d) {
const auto x = 30 * i + residues[d];
pattern[i] |= std::uint8_t(bool(x % 7 && x % 11 && x % 13) << d);
}
return pattern;
}();
std::size_t root = std::sqrt(static_cast<long double>(n - 1));
while (root > (n - 1) / root) --root;
while (root + 1 <= (n - 1) / (root + 1)) ++root;
const auto primes = prime_sieve(root + 1);
struct State {
std::size_t p;
std::array<std::size_t, 8> next;
std::array<std::uint8_t, 8> mask;
};
std::vector<State> state;
state.reserve(primes.size());
for (const auto p : primes) {
if (p < 17) continue;
State current;
current.p = p;
for (unsigned d = 0; d < 8; ++d) {
const std::size_t k = p + (residues[d] + 30 - p % 30) % 30;
const __uint128_t x = __uint128_t(p) * k;
current.next[d] = x / 30;
current.mask[d] = std::uint8_t(~(1u << index[x % 30]));
}
state.push_back(current);
}
const std::size_t size = (n - 1) / 30 + 1;
std::vector<std::uint8_t> prime(block_size);
constexpr auto offsets = [&] {
std::array<unsigned, 64> offsets{};
for (unsigned i = 0; i < 64; ++i) offsets[i] = 30 * (i / 8) + residues[i % 8];
return offsets;
}();
for (std::size_t base = 0; base < size;) {
const std::size_t length = std::min(block_size, size - base);
for (std::size_t i = 0; i < length; i += period)
std::copy_n(pattern.data(), std::min(period, length - i), prime.data() + i);
if (base == 0) {
prime[0] &= ~1u;
prime[0] |= (1u << 1) | (1u << 2) | (1u << 3);
}
for (auto ¤t : state) {
const auto p = current.p;
for (unsigned d = 0; d < 8; ++d) {
std::size_t j = current.next[d];
const auto mask = current.mask[d];
for (; j + 3 * p < length; j += 4 * p) {
prime[j] &= mask;
prime[j + p] &= mask;
prime[j + 2 * p] &= mask;
prime[j + 3 * p] &= mask;
}
for (; j < length; j += p) prime[j] &= mask;
current.next[d] = j - length;
}
}
if (base + length == size) {
const auto remaining = n - 30 * (size - 1);
for (unsigned d = 0; d < 8; ++d) if (residues[d] >= remaining) prime[length - 1] &= ~(1u << d);
}
std::fill(prime.begin() + length, prime.begin() + ((length + 7) / 8 * 8), 0);
for (std::size_t i = 0; i < length; i += 8) {
std::uint64_t word = 0;
std::memcpy(&word, prime.data() + i, sizeof(word));
if constexpr (std::endian::native == std::endian::big) word = std::byteswap(word);
else if constexpr (std::endian::native != std::endian::little) {
word = 0;
for (unsigned d = 0; d < 8; ++d) word |= std::uint64_t(prime[i + d]) << (8 * d);
}
while (word) {
const auto x = 30 * (base + i) + offsets[std::countr_zero(word)];
emit(x);
word &= word - 1;
}
}
base += length;
}
}
inline std::vector<std::size_t> prime_sieve(std::size_t n) {
std::vector<std::size_t> result;
if (n >= 1024) result.reserve(std::size_t(double(n) / (std::log(double(n)) - 1.1)));
for_each_prime(n, [&](std::size_t p) { result.push_back(p); });
return result;
}
#line 7 "components/divisor_zeta.hpp"
namespace divisor_zeta_detail {
template <bool inverse, bool multiples, class F>
void for_each(std::size_t n, const std::vector<std::size_t> &ps, std::size_t count, F update) {
if (!n) return;
for (std::size_t j = 0; j < count; ++j) {
const std::size_t p = ps[j], end = (n - 1) / p;
auto apply = [&](std::size_t i) { update(multiples ? i : i * p, multiples ? i * p : i); };
if constexpr (inverse == multiples) {
std::size_t i = 1;
for (; i + 3 <= end; i += 4) apply(i), apply(i + 1), apply(i + 2), apply(i + 3);
for (; i <= end; ++i) apply(i);
} else {
std::size_t i = end;
for (; i >= 4; i -= 4) apply(i), apply(i - 1), apply(i - 2), apply(i - 3);
for (; i; --i) apply(i);
}
}
}
}
template <bool inverse = false, bool multiples = false, class T>
void divisor_zeta(std::vector<T> &a, const std::vector<std::size_t> &ps) {
divisor_zeta_detail::for_each<inverse, multiples>(a.size(), ps, ps.size(), [&](std::size_t to, std::size_t from) {
if constexpr (inverse) a[to] -= a[from];
else a[to] += a[from];
});
}
template <bool inverse = false, bool multiples = false, class T>
void divisor_zeta(std::vector<T> &a) { divisor_zeta<inverse, multiples>(a, prime_sieve(a.size())); }
#line 10 "components/gcd_convolution_mint.hpp"
#if defined(__x86_64__) || defined(__i386__)
#include <immintrin.h>
#endif
namespace mint_detail { struct Mint; }
namespace gcd_convolution_mint_detail {
#if defined(__x86_64__) || defined(__i386__)
template <bool inverse, class T>
__attribute__((target("avx2")))
void prime_pass(T *a, T *b, std::size_t n, const std::vector<std::size_t> &ps, std::size_t count) {
const __m256i mod = _mm256_set1_epi32(2 * T::mod());
for (std::size_t j = 0; j < count; ++j) {
const std::size_t p = ps[j], end = (n - 1) / p;
auto scalar = [&](std::size_t i) {
if constexpr (inverse) a[i] -= a[i * p];
else a[i] += a[i * p], b[i] += b[i * p];
};
auto gather = [&](const T *s) __attribute__((target("avx2"))) {
return _mm256_setr_epi32(s[0].v, s[p].v, s[2 * p].v, s[3 * p].v,
s[4 * p].v, s[5 * p].v, s[6 * p].v, s[7 * p].v);
};
auto vector = [&](std::size_t i) __attribute__((target("avx2"))) {
__m256i x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + i)), y = gather(a + i * p);
if constexpr (inverse) {
x = _mm256_sub_epi32(x, y);
x = _mm256_min_epu32(x, _mm256_add_epi32(x, mod));
} else {
x = _mm256_add_epi32(x, y);
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, mod));
}
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + i), x);
if constexpr (!inverse) {
x = _mm256_add_epi32(_mm256_loadu_si256(reinterpret_cast<const __m256i *>(b + i)), gather(b + i * p));
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, mod));
_mm256_storeu_si256(reinterpret_cast<__m256i *>(b + i), x);
}
};
if constexpr (inverse) {
std::size_t i = 1;
for (; i <= end && i < 8; ++i) scalar(i);
for (; i + 7 <= end; i += 8) vector(i);
for (; i <= end; ++i) scalar(i);
} else {
std::size_t i = end;
for (; i >= 15; i -= 8) vector(i - 7);
for (; i; --i) scalar(i);
}
}
}
#endif
template <bool inverse, class T>
void transform(std::vector<T> &a, std::vector<T> &b, const std::vector<std::size_t> &ps, std::size_t count) {
#if defined(__x86_64__) || defined(__i386__)
if (__builtin_cpu_supports("avx2")) {
prime_pass<inverse>(a.data(), b.data(), a.size(), ps, count);
return;
}
#endif
divisor_zeta_detail::for_each<inverse, true>(a.size(), ps, count, [&](std::size_t to, std::size_t from) {
if constexpr (inverse) a[to] -= a[from];
else a[to] += a[from], b[to] += b[from];
});
}
template <bool inverse, class T>
void large_multiples(std::vector<T> &a, std::vector<T> &b, const std::vector<std::size_t> &ps,
std::size_t first) {
if (first == ps.size()) return;
constexpr std::uint64_t twice = 2 * T::mod();
const std::size_t n = a.size() - 1;
for (std::size_t i = 1; i <= n / ps[first]; ++i) {
const auto end = std::upper_bound(ps.begin() + first, ps.end(), n / i);
std::uint64_t x = inverse ? 0 : a[i].v, y = inverse ? a[i].v : b[i].v;
for (auto p = ps.begin() + first; p != end; ++p) {
x += a[i * *p].v;
if constexpr (!inverse) y += b[i * *p].v;
}
if constexpr (inverse) a[i] = T::word(y % twice) - T::word(x % twice);
else {
a[i] = T::word(x % twice), b[i] = T::word(y % twice);
}
}
}
template <class T>
void multiply(std::vector<T> &a, std::vector<T> &b, const std::vector<std::size_t> &ps) {
const std::size_t n = a.size() - 1;
std::size_t cutoff = std::sqrt(n);
while (cutoff + 1 <= n / (cutoff + 1)) ++cutoff;
while (cutoff && cutoff > n / cutoff) --cutoff;
const std::size_t first = std::upper_bound(ps.begin(), ps.end(), cutoff) - ps.begin();
transform<false>(a, b, ps, first);
large_multiples<false>(a, b, ps, first);
for (std::size_t i = 1; i <= n; ++i) a[i] *= b[i];
transform<true>(a, b, ps, first);
large_multiples<true>(a, b, ps, first);
}
}
#line 6 "convolution/gcd_lcm_convolution.hpp"
#include <cassert>
#line 9 "convolution/gcd_lcm_convolution.hpp"
#include <type_traits>
#include <utility>
#line 12 "convolution/gcd_lcm_convolution.hpp"
namespace gcd_lcm_convolution_detail {
template <bool gcd, class T>
std::vector<T> multiply(std::vector<T> a, std::vector<T> b) {
if (a.empty() || b.empty()) return {};
assert(a[0] == T(0) && b[0] == T(0));
std::size_t n = std::max(a.size(), b.size());
a.resize(n, T(0));
b.resize(n, T(0));
auto ps = prime_sieve(n);
if constexpr (gcd && std::is_same_v<T, mint_detail::Mint>) {
if (n >= 4096 && n - 1 <= UINT32_MAX) {
gcd_convolution_mint_detail::multiply(a, b, ps);
return a;
}
}
divisor_zeta_detail::for_each<false, gcd>(n, ps, ps.size(), [&](std::size_t to, std::size_t from) {
a[to] += a[from], b[to] += b[from];
});
for (std::size_t i = 1; i < n; ++i) a[i] *= b[i];
divisor_zeta<true, gcd>(a, ps);
return a;
}
}
template <class T>
std::vector<T> gcd_convolution(std::vector<T> a, std::vector<T> b) {
return gcd_lcm_convolution_detail::multiply<true>(std::move(a), std::move(b));
}
template <class T>
std::vector<T> lcm_convolution(std::vector<T> a, std::vector<T> b) {
return gcd_lcm_convolution_detail::multiply<false>(std::move(a), std::move(b));
}