po167_library

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub potato167/po167_library

:heavy_check_mark: 乗法的関数の累積和
(math/multiplicative_prefix_sum.hpp)

概要

乗法的関数 $f$($f(1)=1$)について、

\[F(N)=\sum_{n=1}^{N}f(n)\]

を求める。

次の2つを用意し、sum に渡す。

prime_sum は、素数和の表を作るための補助関数である。

関数一覧

構築

po167::multiplicative_prefix_sum<T> sieve(N);

T には atcoder::modint998244353 などを使う。 内部で $2,3,4,5$ の逆元を使うため、long long などの整数型は使えない。

quotients には、集合

\[Q_N=\left\{\left\lfloor\frac{N}{i}\right\rfloor:1\le i\le N\right\}\]

の要素が降順・重複なしで入る。以下の返り値の配列も、この順序で並ぶ。

prime_sum(value, prefix)

auto P = sieve.prime_sum(value, prefix);

補助関数 $g$ に対して、次の2つの関数を渡す。

引数 返す値
value(p) $g(p)$
prefix(x) $\sum_{n=1}^{x}g(n)$

$g$ は 完全乗法的、すなわち $g(1)=1$ かつ、すべての正整数 $a,b$ に対して $g(ab)=g(a)g(b)$ を満たす必要がある。

返り値は、

\[P[i]=\sum_{\substack{p\le \mathrm{quotients}[i]\\p:\mathrm{prime}}}g(p).\]

注 : 第2引数で渡すのは全整数の和、返ってくるのは素数だけの和である。

代表的な使い方は次のとおり。

$g(n)$ value(p) prefix(x) 得られる表
$1$ $1$ $x$ 素数の個数 $P_0(x)$
$n$ $p$ $x(x+1)/2$ 素数の値の和 $P_1(x)$

求めたい関数が素数上で $f(p)=a+bp$ なら、$P=aP_0+bP_1$ として素数和の表を作れる。

sum(P, prime_power)

auto F = sieve.sum(P, prime_power);
引数 内容
P[i] $\sum_{p\le \mathrm{quotients}[i],\ p:\mathrm{prime}}f(p)$
prime_power(p, e, pe) $f(p^e)$ を返す関数。pe には計算済みの $p^e$ が渡される

返り値は、

\[F[i]=\sum_{n=1}^{\mathrm{quotients}[i]}f(n).\]

$f(1)=1$ は内部で処理するため、prime_power で $e=0$ を扱う必要はない。

get(table, x)

T answer = sieve.get(F, N);

表から上限 $x$ に対応する値を取得する。$x$ は $Q_N$ の要素または $0$ に限る。 get(table, 0) は $0$ を返す。

使用例

Sum of Multiplicative Function では、

\[f(p^e)=ae+bp\]

が与えられる。

素数上では $f(p)=a+bp$ なので、$g(n)=1$ と $g(n)=n$ から素数和を作る。 次は test/verify/ に置く提出用コードの例。

#define PROBLEM "https://judge.yosupo.jp/problem/sum_of_multiplicative_function"

#include "../../math/multiplicative_prefix_sum.hpp"
#include <atcoder/modint>
#include <iostream>

int main(){
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    using mint = atcoder::static_modint<469762049>;
    const mint inv2 = mint(1) / 2;

    int T;
    std::cin >> T;
    while (T--){
        long long N;
        int a, b;
        std::cin >> N >> a >> b;

        po167::multiplicative_prefix_sum<mint> sieve(N);

        // P0(x) = sum_{prime p <= x} 1
        auto P0 = sieve.prime_sum(
            [](long long){ return mint(1); },
            [](long long x){ return mint(x); }
        );

        // P1(x) = sum_{prime p <= x} p
        auto P1 = sieve.prime_sum(
            [](long long p){ return mint(p); },
            [inv2](long long x){
                return mint(x) * (mint(x) + 1) * inv2;
            }
        );

        // sum_{prime p <= x} f(p) = a * P0(x) + b * P1(x)
        for (std::size_t i = 0; i < P1.size(); i++){
            P1[i] = P0[i] * a + P1[i] * b;
        }

        // f(p^e) = a * e + b * p
        auto F = sieve.sum(P1, [a, b](long long p, int e, long long){
            return mint(a) * e + mint(b) * p;
        });
        std::cout << sieve.get(F, N).val() << "\n";
    }
}

計算量

渡す関数の評価と T の演算が $O(1)$ のとき、prime_sum と sum はそれぞれ次の計算量で動作する。

prime_sum を $k$ 回使う場合、全体の時間計算量は $O((k+1)N^{2/3}/\log N)$。

Verified with

Code

#pragma once
#include <cassert>
#include <cmath>
#include <vector>

namespace po167 {

// f(1) = 1, gcd(a, b) = 1 => f(ab) = f(a)f(b).
// result[i] = sum_{1 <= k <= quotients[i]} f(k).
// quotients = {N / i : 1 <= i <= N}, in decreasing order.
// T: a commutative ring in which 2, 3, 4, 5 are invertible.
// Each prime_sum / sum: O(N^(2/3) / log N) time, O(sqrt N) space.
template <class T>
class multiplicative_prefix_sum {
    long long n, root, cut = 0, boundary;
    std::vector<long long> primes;
    T inverse[6];

    static long long isqrt(long long x) {
        long long r = std::sqrt((long double)x);
        while (r && r > x / r) --r;
        while (r + 1 <= x / (r + 1)) ++r;
        return r;
    }

    // Index of the smallest quotient >= x; also accepts non-quotients.
    std::size_t index(long long x) const {
        return x <= root ? quotients.size() - x : n / x - 1;
    }

    // a at floor(quotients[i] / d). The caller guarantees d <= quotients[i].
    T divided(const std::vector<T>& a, std::size_t i, long long d) const {
        long long x = quotients[i] / d;
        return a[x <= root ? quotients.size() - x : (i + 1) * d - 1];
    }

    void add_difference(std::vector<T>& a, const std::vector<T>& d) const {
        T s = 0;
        for (std::size_t i = a.size(); i-- > 0;) {
            s += d[i];
            a[i] += s;
        }
    }

    // Dirichlet convolution of two functions supported on integers > 1
    // whose prime factors are all greater than cut.
    std::vector<T> convolution(const std::vector<T>& a,
                               const std::vector<T>& b) const {
        struct term { long long x; T value; };
        std::vector<term> as, bs;
        const std::size_t m = a.size();
        for (long long x = cut + 1; x <= root; ++x) {
            T u = a[m - x] - a[m - x + 1];
            T v = b[m - x] - b[m - x + 1];
            if (u != T(0)) as.push_back({x, u});
            if (v != T(0)) bs.push_back({x, v});
        }
        std::vector<T> c(m);
        if (as.empty() && bs.empty()) return c;

        // Below boundary, both factors necessarily lie in [1, root].
        for (const auto& u : as) {
            long long limit = boundary / u.x;
            for (const auto& v : bs) {
                if (v.x > limit) break;
                c[index(u.x * v.x)] += u.value * v.value;
            }
        }
        for (std::size_t i = m; i-- > 1;) c[i - 1] += c[i];

        // Above boundary, use the Dirichlet hyperbola formula.
        for (std::size_t i = 0; i < m && quotients[i] > boundary; ++i) {
            long long r = isqrt(quotients[i]);
            T value = -a[m - r] * b[m - r];
            for (const auto& u : as) {
                if (u.x > r) break;
                value += u.value * divided(b, i, u.x);
            }
            for (const auto& v : bs) {
                if (v.x > r) break;
                value += v.value * divided(a, i, v.x);
            }
            c[i] = value;
        }
        return c;
    }

    // exp(a) or log(1 + a), using Dirichlet multiplication.
    // Six factors with all prime factors > cut have product > n.
    std::vector<T> series(const std::vector<T>& a, bool exponential) const {
        std::vector<T> result = a, power = a;
        T coefficient = 1;
        for (int k = 2; k <= 5; ++k) {
            power = convolution(power, a);
            if (exponential) coefficient *= inverse[k];
            else coefficient = k % 2 ? inverse[k] : -inverse[k];
            for (std::size_t i = 0; i < a.size(); ++i)
                result[i] += coefficient * power[i];
        }
        if (exponential) for (T& x : result) x += T(1);
        return result;
    }

public:
    std::vector<long long> quotients;

    explicit multiplicative_prefix_sum(long long N) : n(N), root(0) {
        assert(n >= 0);
        root = isqrt(n);
        while ((cut + 1) * (cut + 1) * (cut + 1) <= root) ++cut;
        boundary = root * (cut + 1);
        for (int k = 1; k <= 5; ++k) inverse[k] = T(1) / T(k);
        for (long long l = 1; l <= n;) {
            long long x = n / l;
            quotients.push_back(x);
            if (x == 1) break;
            l = n / x + 1;
        }
        std::vector<bool> composite(root + 1);
        for (long long p = 2; p <= root; ++p) {
            if (composite[p]) continue;
            primes.push_back(p);
            if (p <= root / p)
                for (long long q = p * p; q <= root; q += p) composite[q] = true;
        }
    }

    // x must be 0 or an element of quotients.
    T get(const std::vector<T>& a, long long x) const {
        assert(a.size() == quotients.size() && 0 <= x && x <= n);
        if (x == 0) return T(0);
        std::size_t i = index(x);
        assert(quotients[i] == x);
        return a[i];
    }

    // g(ab) = g(a)g(b) for ALL a, b, and g(1) = 1.
    // value(p) = g(p), prefix(x) = sum_{1 <= k <= x} g(k), both O(1).
    // Returns sum_{prime p <= quotients[i]} g(p).
    template <class Value, class Prefix>
    std::vector<T> prime_sum(Value value, Prefix prefix) const {
        std::vector<T> a;
        a.reserve(quotients.size());
        for (long long x : quotients) a.push_back(prefix(x));
        for (long long p : primes) {
            if (p > cut) break;
            T fp = value(p);
            for (std::size_t i = 0; i < a.size() && quotients[i] >= p; ++i)
                a[i] -= fp * divided(a, i, p);
        }
        for (T& x : a) x -= T(1);
        a = series(a, false);

        std::vector<T> d(a.size());
        for (long long p : primes) {
            T fp = value(p), power = fp;
            if (p <= cut) {
                d[index(p)] += fp;
                continue;
            }
            long long q = p;
            for (int e = 2; q <= n / p; ++e) {
                q *= p;
                power *= fp;
                d[index(q)] -= power * inverse[e];
            }
        }
        add_difference(a, d);
        return a;
    }

    // prime[i] = sum_{prime p <= quotients[i]} f(p).
    // prime_power(p, e, pe) = f(p^e), where pe = p^e, in O(1).
    // Returns sum_{1 <= k <= quotients[i]} f(k).
    template <class PrimePower>
    std::vector<T> sum(std::vector<T> prime, PrimePower prime_power) const {
        assert(prime.size() == quotients.size());
        if (n == 0) return prime;
        T small = get(prime, cut);
        for (std::size_t i = 0; i < prime.size(); ++i)
            prime[i] = quotients[i] > cut ? prime[i] - small : T(0);

        std::vector<T> d(prime.size());
        for (long long p : primes) {
            if (p <= cut) continue;
            T a[6], logarithm[6];
            a[1] = logarithm[1] = prime_power(p, 1, p);
            long long q = p;
            for (int e = 2; q <= n / p; ++e) {
                q *= p;
                a[e] = prime_power(p, e, q);
                T s = 0;
                for (int j = 1; j < e; ++j) s += T(j) * logarithm[j] * a[e - j];
                logarithm[e] = a[e] - s * inverse[e];
                d[index(q)] += logarithm[e];
            }
        }
        add_difference(prime, d);
        prime = series(prime, true);

        for (long long p : primes) {
            if (p > cut) break;
            std::vector<long long> powers;
            std::vector<T> coefficients;
            for (long long q = p, e = 1;; q *= p, ++e) {
                powers.push_back(q);
                coefficients.push_back(prime_power(p, (int)e, q));
                if (q > n / p) break;
            }
            for (std::size_t i = 0; i < prime.size() && quotients[i] >= p; ++i)
                for (std::size_t j = 0; j < powers.size() && powers[j] <= quotients[i]; ++j)
                    prime[i] += coefficients[j] * divided(prime, i, powers[j]);
        }
        return prime;
    }
};

} // namespace po167
#line 2 "math/multiplicative_prefix_sum.hpp"
#include <cassert>
#include <cmath>
#include <vector>

namespace po167 {

// f(1) = 1, gcd(a, b) = 1 => f(ab) = f(a)f(b).
// result[i] = sum_{1 <= k <= quotients[i]} f(k).
// quotients = {N / i : 1 <= i <= N}, in decreasing order.
// T: a commutative ring in which 2, 3, 4, 5 are invertible.
// Each prime_sum / sum: O(N^(2/3) / log N) time, O(sqrt N) space.
template <class T>
class multiplicative_prefix_sum {
    long long n, root, cut = 0, boundary;
    std::vector<long long> primes;
    T inverse[6];

    static long long isqrt(long long x) {
        long long r = std::sqrt((long double)x);
        while (r && r > x / r) --r;
        while (r + 1 <= x / (r + 1)) ++r;
        return r;
    }

    // Index of the smallest quotient >= x; also accepts non-quotients.
    std::size_t index(long long x) const {
        return x <= root ? quotients.size() - x : n / x - 1;
    }

    // a at floor(quotients[i] / d). The caller guarantees d <= quotients[i].
    T divided(const std::vector<T>& a, std::size_t i, long long d) const {
        long long x = quotients[i] / d;
        return a[x <= root ? quotients.size() - x : (i + 1) * d - 1];
    }

    void add_difference(std::vector<T>& a, const std::vector<T>& d) const {
        T s = 0;
        for (std::size_t i = a.size(); i-- > 0;) {
            s += d[i];
            a[i] += s;
        }
    }

    // Dirichlet convolution of two functions supported on integers > 1
    // whose prime factors are all greater than cut.
    std::vector<T> convolution(const std::vector<T>& a,
                               const std::vector<T>& b) const {
        struct term { long long x; T value; };
        std::vector<term> as, bs;
        const std::size_t m = a.size();
        for (long long x = cut + 1; x <= root; ++x) {
            T u = a[m - x] - a[m - x + 1];
            T v = b[m - x] - b[m - x + 1];
            if (u != T(0)) as.push_back({x, u});
            if (v != T(0)) bs.push_back({x, v});
        }
        std::vector<T> c(m);
        if (as.empty() && bs.empty()) return c;

        // Below boundary, both factors necessarily lie in [1, root].
        for (const auto& u : as) {
            long long limit = boundary / u.x;
            for (const auto& v : bs) {
                if (v.x > limit) break;
                c[index(u.x * v.x)] += u.value * v.value;
            }
        }
        for (std::size_t i = m; i-- > 1;) c[i - 1] += c[i];

        // Above boundary, use the Dirichlet hyperbola formula.
        for (std::size_t i = 0; i < m && quotients[i] > boundary; ++i) {
            long long r = isqrt(quotients[i]);
            T value = -a[m - r] * b[m - r];
            for (const auto& u : as) {
                if (u.x > r) break;
                value += u.value * divided(b, i, u.x);
            }
            for (const auto& v : bs) {
                if (v.x > r) break;
                value += v.value * divided(a, i, v.x);
            }
            c[i] = value;
        }
        return c;
    }

    // exp(a) or log(1 + a), using Dirichlet multiplication.
    // Six factors with all prime factors > cut have product > n.
    std::vector<T> series(const std::vector<T>& a, bool exponential) const {
        std::vector<T> result = a, power = a;
        T coefficient = 1;
        for (int k = 2; k <= 5; ++k) {
            power = convolution(power, a);
            if (exponential) coefficient *= inverse[k];
            else coefficient = k % 2 ? inverse[k] : -inverse[k];
            for (std::size_t i = 0; i < a.size(); ++i)
                result[i] += coefficient * power[i];
        }
        if (exponential) for (T& x : result) x += T(1);
        return result;
    }

public:
    std::vector<long long> quotients;

    explicit multiplicative_prefix_sum(long long N) : n(N), root(0) {
        assert(n >= 0);
        root = isqrt(n);
        while ((cut + 1) * (cut + 1) * (cut + 1) <= root) ++cut;
        boundary = root * (cut + 1);
        for (int k = 1; k <= 5; ++k) inverse[k] = T(1) / T(k);
        for (long long l = 1; l <= n;) {
            long long x = n / l;
            quotients.push_back(x);
            if (x == 1) break;
            l = n / x + 1;
        }
        std::vector<bool> composite(root + 1);
        for (long long p = 2; p <= root; ++p) {
            if (composite[p]) continue;
            primes.push_back(p);
            if (p <= root / p)
                for (long long q = p * p; q <= root; q += p) composite[q] = true;
        }
    }

    // x must be 0 or an element of quotients.
    T get(const std::vector<T>& a, long long x) const {
        assert(a.size() == quotients.size() && 0 <= x && x <= n);
        if (x == 0) return T(0);
        std::size_t i = index(x);
        assert(quotients[i] == x);
        return a[i];
    }

    // g(ab) = g(a)g(b) for ALL a, b, and g(1) = 1.
    // value(p) = g(p), prefix(x) = sum_{1 <= k <= x} g(k), both O(1).
    // Returns sum_{prime p <= quotients[i]} g(p).
    template <class Value, class Prefix>
    std::vector<T> prime_sum(Value value, Prefix prefix) const {
        std::vector<T> a;
        a.reserve(quotients.size());
        for (long long x : quotients) a.push_back(prefix(x));
        for (long long p : primes) {
            if (p > cut) break;
            T fp = value(p);
            for (std::size_t i = 0; i < a.size() && quotients[i] >= p; ++i)
                a[i] -= fp * divided(a, i, p);
        }
        for (T& x : a) x -= T(1);
        a = series(a, false);

        std::vector<T> d(a.size());
        for (long long p : primes) {
            T fp = value(p), power = fp;
            if (p <= cut) {
                d[index(p)] += fp;
                continue;
            }
            long long q = p;
            for (int e = 2; q <= n / p; ++e) {
                q *= p;
                power *= fp;
                d[index(q)] -= power * inverse[e];
            }
        }
        add_difference(a, d);
        return a;
    }

    // prime[i] = sum_{prime p <= quotients[i]} f(p).
    // prime_power(p, e, pe) = f(p^e), where pe = p^e, in O(1).
    // Returns sum_{1 <= k <= quotients[i]} f(k).
    template <class PrimePower>
    std::vector<T> sum(std::vector<T> prime, PrimePower prime_power) const {
        assert(prime.size() == quotients.size());
        if (n == 0) return prime;
        T small = get(prime, cut);
        for (std::size_t i = 0; i < prime.size(); ++i)
            prime[i] = quotients[i] > cut ? prime[i] - small : T(0);

        std::vector<T> d(prime.size());
        for (long long p : primes) {
            if (p <= cut) continue;
            T a[6], logarithm[6];
            a[1] = logarithm[1] = prime_power(p, 1, p);
            long long q = p;
            for (int e = 2; q <= n / p; ++e) {
                q *= p;
                a[e] = prime_power(p, e, q);
                T s = 0;
                for (int j = 1; j < e; ++j) s += T(j) * logarithm[j] * a[e - j];
                logarithm[e] = a[e] - s * inverse[e];
                d[index(q)] += logarithm[e];
            }
        }
        add_difference(prime, d);
        prime = series(prime, true);

        for (long long p : primes) {
            if (p > cut) break;
            std::vector<long long> powers;
            std::vector<T> coefficients;
            for (long long q = p, e = 1;; q *= p, ++e) {
                powers.push_back(q);
                coefficients.push_back(prime_power(p, (int)e, q));
                if (q > n / p) break;
            }
            for (std::size_t i = 0; i < prime.size() && quotients[i] >= p; ++i)
                for (std::size_t j = 0; j < powers.size() && powers[j] <= quotients[i]; ++j)
                    prime[i] += coefficients[j] * divided(prime, i, powers[j]);
        }
        return prime;
    }
};

} // namespace po167
Back to top page