math/multiplicative_function/prefix_sum.hpp

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

Depends on

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

Verified with

• verify/yosupo_math/prefix_sum_of_multiplicative_function.test.cpp

Code

#ifndef KK2_MATH_MULTIPLICATIVE_FUNCTION_PREFIX_SUM_HPP
#define KK2_MATH_MULTIPLICATIVE_FUNCTION_PREFIX_SUM_HPP 1

#include <algorithm>
#include <vector>

#include "../enumerate_quotients.hpp"
#include "../prime_table.hpp"

namespace kk2 {

template <class T> struct PrefixSumOfMultiplicativeFunction {
    long long n;
    EnumerateQuotients<long long> eq;
    std::vector<T> prefix_sum_only_prime;
    std::vector<T> prefix_sum;

    PrefixSumOfMultiplicativeFunction(long long n)
        : n(n),
          eq(n),
          prefix_sum_only_prime(eq.size()),
          prefix_sum(eq.size()) {}

    int size() const { return eq.size(); }

    template <T (*f)(long long)> void LucyDP(std::vector<T> &dp, T a) {
        LucyDP([](long long x) { return f(x); }, dp);
    }

    // f is completely multiplicative function
    template <class F> void LucyDP(const F &f, std::vector<T> &dp) {
        assert((int)dp.size() == eq.size());
        PrimeTable::set_upper(eq.sqrt_n);
        for (const long long p : PrimeTable::primes()) {
            if (p > eq.sqrt_n) break;
            T fp = f(p);
            for (int i = eq.size() - 1;; --i) {
                if (eq[i] < p * p) break;
                dp[i] -= (dp[eq.idx(eq[i] / p)] - dp[p - 2]) * fp;
            }
        }
    }

    template <T (*f)(long long, long long)> void Min_25Sieve() {
        Min_25Sieve([](long long x, long long y) { return f(x, y); });
    }

    // f is multiplicative function
    template <class F> void Min_25Sieve(const F &f) {
        PrimeTable::set_upper(eq.sqrt_n);
        std::copy(prefix_sum_only_prime.begin(), prefix_sum_only_prime.end(), prefix_sum.begin());
        const auto &primes = PrimeTable::primes();
        std::vector<T> tmp(eq.size());
        for (int i = std::upper_bound(primes.begin(), primes.end(), eq.sqrt_n) - primes.begin() - 1;
             i >= 0;
             --i) {
            const long long p = primes[i];
            T pk = f(p, 1);
            T pk1;
            for (long long p_pw = p, k = 1; n / p >= p_pw; ++k, p_pw *= p) {
                T pk1 = f(p, k + 1);
                for (int j = eq.size() - 1;; --j) {
                    if (eq[j] < p_pw * p) break;
                    tmp[j] += pk * (prefix_sum[eq.idx(eq[j] / p_pw)] - prefix_sum_only_prime[p - 1])
                              + pk1;
                }
                pk = pk1;
            }
            for (int j = eq.size() - 1;; --j) {
                if (eq[j] < p * p) break;
                prefix_sum[j] += tmp[j];
                tmp[j] = T();
            }
        }
        for (int i = 0; i < eq.size(); ++i) ++prefix_sum[i];
    }
};

} // namespace kk2

#endif // KK2_MATH_MULTIPLICATIVE_FUNCTION_PREFIX_SUM_HPP

#line 1 "math/multiplicative_function/prefix_sum.hpp"



#include <algorithm>
#include <vector>

#line 1 "math/enumerate_quotients.hpp"



#include <numeric>
#line 6 "math/enumerate_quotients.hpp"

#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/enumerate_quotients.hpp"

namespace kk2 {

template <class T> struct EnumerateQuotients {
    T n;
    int sqrt_n;
    std::vector<T> res;

    EnumerateQuotients(T n) : n(n), sqrt_n(sqrt_floor(n)) {
        res.resize(sqrt_n + n / (sqrt_n + 1));
        std::iota(res.begin(), res.begin() + sqrt_n, 1);
        for (T i = n / (sqrt_n + 1), j = sqrt_n; i; --i, ++j) res[j] = n / i;
    }

    const std::vector<T> &get() const { return res; }

    int size() const { return res.size(); }

    const T &operator[](int i) const { return res[i]; }

    int idx(T x) const {
        if (x <= sqrt_n) return x - 1;
        return size() - n / x;
    }
};

} // namespace kk2


#line 1 "math/prime_table.hpp"



#line 6 "math/prime_table.hpp"

#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 9 "math/multiplicative_function/prefix_sum.hpp"

namespace kk2 {

template <class T> struct PrefixSumOfMultiplicativeFunction {
    long long n;
    EnumerateQuotients<long long> eq;
    std::vector<T> prefix_sum_only_prime;
    std::vector<T> prefix_sum;

    PrefixSumOfMultiplicativeFunction(long long n)
        : n(n),
          eq(n),
          prefix_sum_only_prime(eq.size()),
          prefix_sum(eq.size()) {}

    int size() const { return eq.size(); }

    template <T (*f)(long long)> void LucyDP(std::vector<T> &dp, T a) {
        LucyDP([](long long x) { return f(x); }, dp);
    }

    // f is completely multiplicative function
    template <class F> void LucyDP(const F &f, std::vector<T> &dp) {
        assert((int)dp.size() == eq.size());
        PrimeTable::set_upper(eq.sqrt_n);
        for (const long long p : PrimeTable::primes()) {
            if (p > eq.sqrt_n) break;
            T fp = f(p);
            for (int i = eq.size() - 1;; --i) {
                if (eq[i] < p * p) break;
                dp[i] -= (dp[eq.idx(eq[i] / p)] - dp[p - 2]) * fp;
            }
        }
    }

    template <T (*f)(long long, long long)> void Min_25Sieve() {
        Min_25Sieve([](long long x, long long y) { return f(x, y); });
    }

    // f is multiplicative function
    template <class F> void Min_25Sieve(const F &f) {
        PrimeTable::set_upper(eq.sqrt_n);
        std::copy(prefix_sum_only_prime.begin(), prefix_sum_only_prime.end(), prefix_sum.begin());
        const auto &primes = PrimeTable::primes();
        std::vector<T> tmp(eq.size());
        for (int i = std::upper_bound(primes.begin(), primes.end(), eq.sqrt_n) - primes.begin() - 1;
             i >= 0;
             --i) {
            const long long p = primes[i];
            T pk = f(p, 1);
            T pk1;
            for (long long p_pw = p, k = 1; n / p >= p_pw; ++k, p_pw *= p) {
                T pk1 = f(p, k + 1);
                for (int j = eq.size() - 1;; --j) {
                    if (eq[j] < p_pw * p) break;
                    tmp[j] += pk * (prefix_sum[eq.idx(eq[j] / p_pw)] - prefix_sum_only_prime[p - 1])
                              + pk1;
                }
                pk = pk1;
            }
            for (int j = eq.size() - 1;; --j) {
                if (eq[j] < p * p) break;
                prefix_sum[j] += tmp[j];
                tmp[j] = T();
            }
        }
        for (int i = 0; i < eq.size(); ++i) ++prefix_sum[i];
    }
};

} // namespace kk2

Back to top page