convolution/divisor_multiple_transform.hpp

• View this file on GitHub
• Last update: 2025-04-24 20:54:26+09:00
• Include: #include "convolution/divisor_multiple_transform.hpp"

Depends on

• math/frac_floor.hpp
• math/prime_table.hpp
• math/sqrt_floor.hpp

Required by

• convolution/convolution_gcd.hpp
• convolution/convolution_lcm.hpp

Verified with

• verify/yosupo_convolution/convolution_gcd.test.cpp
• verify/yosupo_convolution/convolution_lcm.test.cpp

Code

#ifndef KK2_CONVLUTION_DIVISOR_MULTIPLE_TRANSFORM_HPP
#define KK2_CONVLUTION_DIVISOR_MULTIPLE_TRANSFORM_HPP 1

#include "../math/prime_table.hpp"

namespace kk2 {

template <class FPS> void multiple_transform(FPS &a) {
    int n = int(a.size());
    if (!n) return;
    n--;
    PrimeTable::set_upper(n);
    for (const auto p : PrimeTable::primes()) {
        if (p > n) break;
        for (int i = n / p; i; i--) a[i] += a[i * p];
    }
}

template <class FPS> void inverse_multiple_transform(FPS &a) {
    int n = int(a.size());
    if (!n) return;
    n--;
    PrimeTable::set_upper(n);
    for (const auto p : PrimeTable::primes()) {
        if (p > n) break;
        for (int i = 1; i <= n / p; i++) a[i] -= a[i * p];
    }
}

template <class FPS> void divisor_transform(FPS &a) {
    int n = int(a.size());
    if (!n) return;
    n--;
    PrimeTable::set_upper(n);
    for (const auto p : PrimeTable::primes()) {
        if (p > n) break;
        for (int i = 1; i <= n / p; i++) a[i * p] += a[i];
    }
}

template <class FPS> void inverse_divisor_transform(FPS &a) {
    int n = int(a.size());
    if (!n) return;
    n--;
    PrimeTable::set_upper(n);
    for (const auto p : PrimeTable::primes()) {
        if (p > n) break;
        for (int i = n / p; i > 0; i--) a[i * p] -= a[i];
    }
}

} // namespace kk2

#endif // KK2_CONVLUTION_DIVISOR_MULTIPLE_TRANSFORM_HPP

#line 1 "convolution/divisor_multiple_transform.hpp"



#line 1 "math/prime_table.hpp"



#include <algorithm>
#include <vector>

#line 1 "math/sqrt_floor.hpp"



#include <cmath>

#line 1 "math/frac_floor.hpp"



#include <cassert>

namespace kk2 {

// floor(x) = ceil(x) - 1 (for all x not in Z) ...(1)
// floor(x) = -ceil(-x)   (for all x)          ...(2)

// return floor(a / b)
template <typename T, typename U> constexpr T fracfloor(T a, U b) {
    assert(b != 0);
    if (a % b == 0) return a / b;
    if (a >= 0) return a / b;

    // floor(x) = -ceil(-x)      by (2)
    //          = -floor(-x) - 1 by (1)
    return -((-a) / b) - 1;
}

// return ceil(a / b)
template <typename T, typename U> constexpr T fracceil(T a, U b) {
    assert(b != 0);
    if (a % b == 0) return a / b;
    if (a >= 0) return a / b + 1;

    // ceil(x) = -floor(-x)      by (2)
    return -((-a) / b);
}

} // namespace kk2


#line 7 "math/sqrt_floor.hpp"

namespace kk2 {

template <typename T> T sqrt_floor(T n) {
    assert(n >= 0);
    if (n == T(0)) return 0;
    T x = std::sqrt(n);
    if (x == T(0)) ++x;
    while (x > kk2::fracfloor(n, x)) --x;
    while (x + 1 <= kk2::fracfloor(n, x + 1)) ++x;
    return x;
}

template <typename T> T sqrt_ceil(T n) {
    assert(n >= 0);
    if (n <= T(1)) return n;
    T x = std::sqrt(n);
    if (x == T(0)) ++x;
    while (x < kk2::fracceil(n, x)) ++x;
    while (x - 1 >= kk2::fracceil(n, x - 1)) --x;
    return x;
}

} // namespace kk2


#line 8 "math/prime_table.hpp"

namespace kk2 {

struct PrimeTable {
  private:
    static inline int _n = 30;
    static inline std::vector<int> _primes{2, 3, 5, 7, 11, 13, 17, 19, 23, 29};

  public:
    PrimeTable() = delete;

    // wheel sieve
    // reference: https://37zigen.com/wheel-sieve/
    static void set_upper(int m, int reserve_size = 26355867) {
        if (m <= _n) return;
        _n = std::max(m, 2 * _n);
        int sqrt_n = sqrt_floor(_n);
        int w = 1;
        std::vector<bool> iscoprime(sqrt_n, true);
        for (int i = 0; i < 9; i++) {
            if (w * _primes[i] > sqrt_n) break;
            w *= _primes[i];
            for (int j = _primes[i]; j < sqrt_n; j += _primes[i]) iscoprime[j] = false;
        }

        std::vector<int> idx_(w, -1);
        int s = 0;
        for (int i = 1; i < w; i++) {
            if (iscoprime[i]) idx_[i] = s++;
        }
        std::vector<int> coprimes(s);
        for (int i = 1; i < w; i++) {
            if (idx_[i] != -1) coprimes[idx_[i]] = i;
        }

        auto idx = [&](long long x) -> long long {
            if (idx_[x % w] == -1) return -1;
            return x / w * s + idx_[x % w];
        };

        auto val = [&](int i) {
            return i / s * w + coprimes[i % s];
        };

        int n = (_n + w - 1) / w * s;
        std::vector<int> _primes2;
        _primes2.reserve(reserve_size);
        std::vector<int> lpf(n, 0);
        for (int i = 1; i < n; i++) {
            int v = val(i);
            if (lpf[i] == 0) {
                lpf[i] = v;
                _primes2.push_back(lpf[i]);
            }

            for (const long long p : _primes2) {
                long long j = idx(p * v);
                if (j >= n) break;
                if (lpf[i] < p) break;
                lpf[j] = p;
            }
        }

        std::vector<int> tmp;
        tmp.reserve(_primes.size() + _primes2.size());
        std::set_union(_primes.begin(),
                       _primes.end(),
                       _primes2.begin(),
                       _primes2.end(),
                       std::back_inserter(tmp));
        _primes = std::move(tmp);
    }

    static const std::vector<int> &primes() { return _primes; }

    template <typename It> struct PrimeIt {
        It bg, ed;
        PrimeIt(It bg_, It ed_) : bg(bg_), ed(ed_) {}
        It begin() const { return bg; }
        It end() const { return ed; }
        int size() const { return ed - bg; }
        int operator[](int i) const { return bg[i]; }
        std::vector<int> to_vec() const { return std::vector<int>(bg, ed); }
    };

    static auto primes(int n) {
        if (n >= _n) set_upper(n);
        return PrimeIt(_primes.begin(), std::upper_bound(_primes.begin(), _primes.end(), n));
    }
};

} // namespace kk2


#line 5 "convolution/divisor_multiple_transform.hpp"

namespace kk2 {

template <class FPS> void multiple_transform(FPS &a) {
    int n = int(a.size());
    if (!n) return;
    n--;
    PrimeTable::set_upper(n);
    for (const auto p : PrimeTable::primes()) {
        if (p > n) break;
        for (int i = n / p; i; i--) a[i] += a[i * p];
    }
}

template <class FPS> void inverse_multiple_transform(FPS &a) {
    int n = int(a.size());
    if (!n) return;
    n--;
    PrimeTable::set_upper(n);
    for (const auto p : PrimeTable::primes()) {
        if (p > n) break;
        for (int i = 1; i <= n / p; i++) a[i] -= a[i * p];
    }
}

template <class FPS> void divisor_transform(FPS &a) {
    int n = int(a.size());
    if (!n) return;
    n--;
    PrimeTable::set_upper(n);
    for (const auto p : PrimeTable::primes()) {
        if (p > n) break;
        for (int i = 1; i <= n / p; i++) a[i * p] += a[i];
    }
}

template <class FPS> void inverse_divisor_transform(FPS &a) {
    int n = int(a.size());
    if (!n) return;
    n--;
    PrimeTable::set_upper(n);
    for (const auto p : PrimeTable::primes()) {
        if (p > n) break;
        for (int i = n / p; i > 0; i--) a[i * p] -= a[i];
    }
}

} // namespace kk2

Back to top page