Dirichlet Series
(math/DirichletSeries.hpp)

• View this file on GitHub
• Last update: 2024-05-21 15:55:19+09:00
• Include: #include "math/DirichletSeries.hpp"

説明

Dirichlet級数を取り扱う構造体。 $a$ に関するDirichlet級数 $D_a(s)$ は

\[D_a(s) = \sum_{n = 1}^{\infty} a_n s^{-n}\]

で表される。これを長さ $N$ で打ち切ったものについて取り扱う。

Dirichlet級数についてと実装方法について、maspyさんの記事を参考にした。

コンストラクタ

$a_i$ を返す関数 $f$ とその累積和 $A_i$ を計算する関数 $F$ を引数として渡すことで、 $a$ に関するDirichlet級数を生成する。

set_size(n)をすることを忘れずにすること。

和・差

$O(K + L)$

積

$O(K\log K + (NL)^{1/2})$

商

$O(K\log K + (NL)^{1/2})$

get_sum()

$\sum_{n = 1}^{N} a_n$ を返す。 $O(1)$

pow(u64 n)

Dirichlet級数 $a$ について $a^n$ を求める。繰り返し二乗法で $\log n$ 回程度の積が実行される。

zeta()

$\zeta(s) = \sum_{n = 1}^{\infty} s^{-n}$ のDirichlet級数を返す。

zeta1()

$\zeta(s-1) = \sum_{n = 1}^{\infty} n \times s^{-n}$ のDirichlet級数を返す。

set_size(i64 n)

長さを $n$ にする。

set_size_multiplicative(i64 n)

長さを $n$ として、 $a$ が乗法的関数のとき、計算量が良くなるような分割になるようにする。

Depends on

• Dirichlet Convolution (convolution/dirichlet_convolution.hpp)
• Eratosthenes Sieve (math/eratosthenes_sieve.hpp)
• Linear Sieve (math/linear_sieve.hpp)
• modint/base.hpp
• template/int_alias.hpp

Verified with

• test/math/Sum_of_Totient_Function.test.cpp

Code

#pragma once

#include <functional>

#include <numeric>

#include <vector>


#include "../convolution/dirichlet_convolution.hpp"

#include "../template/int_alias.hpp"


namespace ebi {

template <class T, int id> struct DirichletSeries {
  private:
    using Self = DirichletSeries<T, id>;

    void set(std::function<T(i64)> f, std::function<T(i64)> F) {
        for (int i = 1; i <= K; i++) {
            a[i] = f(i);
        }
        for (int i = 1; i <= L; i++) {
            A[i] = F(N / i);
        }
    }

  public:
    DirichletSeries() : a(K + 1), A(L + 1) {}

    DirichletSeries(std::function<T(i64)> f, std::function<T(i64)> F,
                    bool _is_multiplicative = false)
        : a(K + 1), A(L + 1), is_multiplicative(_is_multiplicative) {
        set(f, F);
    }

    Self operator+(const Self &rhs) const noexcept {
        return Self(*this) += rhs;
    }
    Self operator-(const Self &rhs) const noexcept {
        return Self(*this) -= rhs;
    }
    Self operator*(const Self &rhs) const noexcept {
        return Self(*this) *= rhs;
    }
    Self operator/(const Self &rhs) const noexcept {
        return Self(*this) /= rhs;
    }

    Self operator+=(const Self &rhs) noexcept {
        for (int i = 1; i <= K; i++) {
            a[i] += rhs.a[i];
        }
        for (int i = 1; i <= L; i++) {
            A[i] += rhs.A[i];
        }
        return *this;
    }

    Self operator-=(const Self &rhs) noexcept {
        for (int i = 1; i <= K; i++) {
            a[i] -= rhs.a[i];
        }
        for (int i = 1; i <= L; i++) {
            A[i] -= rhs.A[i];
        }
        return *this;
    }

    Self operator*=(const Self &rhs) noexcept {
        Self ret;
        if (is_multiplicative && rhs.is_multiplicative) {
            ret.a = dirichlet_convolution_multiplicative_function(a, rhs.a);
            ret.is_multiplicative = true;
        } else if (is_multiplicative) {
            ret.a =
                dirichlet_convolution_left_is_multiplicative_function(a, rhs.a);
        } else if (rhs.is_multiplicative) {
            ret.a =
                dirichlet_convolution_left_is_multiplicative_function(rhs.a, a);
        } else {
            ret.a = dirichlet_convolution(a, rhs.a);
        }
        std::vector<T> sum_a = a, sum_b = rhs.a;
        for (int i = 1; i < K; i++) {
            sum_a[i + 1] += sum_a[i];
            sum_b[i + 1] += sum_b[i];
        }
        auto get_A = [&](i64 x) -> T {
            if (x <= K) {
                return sum_a[x];
            } else {
                return A[N / x];
            }
        };
        auto get_B = [&](i64 x) -> T {
            if (x <= K) {
                return sum_b[x];
            } else {
                return rhs.A[N / x];
            }
        };
        for (i64 l = L, m = 1; l >= 1; l--) {
            i64 n = N / l;
            while (m * m <= n) m++;
            m--;
            for (int i = 1; i <= m; i++) {
                ret.A[l] +=
                    a[i] * get_B(n / i) + (get_A(n / i) - get_A(m)) * rhs.a[i];
            }
        }
        return ret;
    }

    // c = a / b

    Self operator/=(const Self &b) noexcept {
        Self c = *this;
        T inv_a = b.a[1].inv();
        for (int i = 1; i <= K; i++) {
            c.a[i] *= inv_a;
            for (int j = 2; i * j <= K; j++) {
                c.a[i * j] -= c.a[i] * b.a[j];
            }
        }
        std::vector<T> sum_b = b.a, sum_c = c.a;
        for (int i = 1; i < K; ++i) {
            sum_b[i + 1] += sum_b[i];
            sum_c[i + 1] += sum_c[i];
        }
        auto get_B = [&](i64 x) -> T {
            if (x <= K) {
                return sum_b[x];
            } else {
                return b.A[N / x];
            }
        };
        auto get_C = [&](i64 x) -> T {
            if (x <= K) {
                return sum_c[x];
            } else {
                return c.A[N / x];
            }
        };
        for (i64 l = L, m = 1; l >= 1; l--) {
            i64 n = N / l;
            while (m * m <= n) m++;
            m--;
            for (int i = 2; i <= m; i++) {
                c.A[l] -= b.a[i] * get_C(n / i);
            }

            for (int i = 1; i <= m; i++) {
                c.A[l] -= c.a[i] * (get_B(n / i) - get_B(m));
            }
            c.A[l] *= inv_a;
        }
        return *this = c;
    }

    Self pow(u64 n) const {
        Self res;
        res.a[1] = 1;
        res.is_multiplicative = is_multiplicative;
        std::fill(res.A.begin(), res.A.end(), 1);
        Self x = *this;
        while (n > 0) {
            if (n & 1) res = res * x;
            x = x * x;
            n >>= 1;
        }
        return res;
    }

    T get_sum() {
        return A[1];
    }

    static Self zeta() {
        Self ret;
        std::fill(ret.a.begin(), ret.a.end(), 1);
        for (int i = 1; i <= L; i++) {
            ret.A[i] = N / i;
        }
        ret.is_multiplicative = true;
        return ret;
    }

    static Self zeta1() {
        Self ret;
        ret.is_multiplicative = true;
        std::iota(ret.a.begin(), ret.a.end(), 0);
        T inv2 = T(2).inv();
        for (int i = 1; i <= L; i++) {
            i64 n = N / i;
            ret.A[i] = T(n) * T(n + 1) * inv2;
        }
        return ret;
    }

    static Self zeta2() {
        Self ret;
        ret.is_multiplicative = true;
        for (int i = 1; i <= K; i++) {
            ret.a[i] = i * i;
        }
        T inv6 = T(6).inv();
        for (int i = 1; i <= L; i++) {
            i64 n = N / i;
            ret.A[i] = T(n) * T(n + 1) * T(2 * n + 1) * inv6;
        }
    }

    static void set_size(i64 n) {
        N = n;
        if (N <= 10) {
            K = N;
            L = 1;
        } else if (N <= 5000) {
            K = 1;
            while (K * K < N) K++;
            L = (N + K - 1) / K;
        } else {
            L = 1;
            while (L * L * L / 50 < N) L++;
            K = (N + L - 1) / L;
        }
    }

    static void set_size_multiplicative(i64 n) {
        N = n;
        L = 1;
        while (L * L * L < N) L++;
        K = L * L;
    }

  private:
    static i64 N, K, L;
    static std::vector<std::pair<int, int>> prime_pow_table;
    std::vector<T> a, A;
    bool is_multiplicative = false;
};

template <class T, int id> i64 DirichletSeries<T, id>::N = 1000000;
template <class T, int id> i64 DirichletSeries<T, id>::K = 10000;
template <class T, int id> i64 DirichletSeries<T, id>::L = 100;
template <class T, int id>
std::vector<std::pair<int, int>> DirichletSeries<T, id>::prime_pow_table = {};

}  // namespace ebi

#line 2 "math/DirichletSeries.hpp"

#include <functional>

#include <numeric>

#include <vector>


#line 2 "convolution/dirichlet_convolution.hpp"

#line 4 "convolution/dirichlet_convolution.hpp"

#line 2 "math/eratosthenes_sieve.hpp"

#include <cassert>

#include <cstdint>

#line 6 "math/eratosthenes_sieve.hpp"

/*
    reference: https://37zigen.com/sieve-eratosthenes/
*/

namespace ebi {

struct eratosthenes_sieve {
  private:
    using i64 = std::int_fast64_t;
    int n;
    std::vector<bool> table;

  public:
    eratosthenes_sieve(int _n) : n(_n), table(std::vector<bool>(n + 1, true)) {
        table[1] = false;
        for (i64 i = 2; i * i <= n; i++) {
            if (!table[i]) continue;
            for (i64 j = i; i * j <= n; j++) {
                table[i * j] = false;
            }
        }
    }

    bool is_prime(int p) {
        return table[p];
    }

    std::vector<int> prime_table(int m = -1) {
        if (m < 0) m = n;
        std::vector<int> prime;
        for (int i = 2; i <= m; i++) {
            if (table[i]) prime.emplace_back(i);
        }
        return prime;
    }
};

}  // namespace ebi
#line 2 "math/linear_sieve.hpp"

#line 2 "modint/base.hpp"

#include <concepts>
#include <iostream>
#include <utility>

namespace ebi {

template <class T>
concept Modint = requires(T a, T b) {
    a + b;
    a - b;
    a * b;
    a / b;
    a.inv();
    a.val();
    a.pow(std::declval<long long>());
    T::mod();
};

template <Modint mint> std::istream &operator>>(std::istream &os, mint &a) {
    long long x;
    os >> x;
    a = x;
    return os;
}

template <Modint mint>
std::ostream &operator<<(std::ostream &os, const mint &a) {
    return os << a.val();
}

}  // namespace ebi
#line 2 "template/int_alias.hpp"

#line 4 "template/int_alias.hpp"

namespace ebi {

using ld = long double;
using std::size_t;
using i8 = std::int8_t;
using u8 = std::uint8_t;
using i16 = std::int16_t;
using u16 = std::uint16_t;
using i32 = std::int32_t;
using u32 = std::uint32_t;
using i64 = std::int64_t;
using u64 = std::uint64_t;
using i128 = __int128_t;
using u128 = __uint128_t;

}  // namespace ebi
#line 5 "math/linear_sieve.hpp"

/*
    reference: https://37zigen.com/linear-sieve/
    verify:    https://atcoder.jp/contests/abc162/submissions/25095562
*/

#line 13 "math/linear_sieve.hpp"

namespace ebi {

struct linear_sieve {
  private:
    using u64 = std::uint64_t;
    int n;
    std::vector<int> sieve;
    std::vector<int> prime;

  public:
    linear_sieve(int _n) : n(_n), sieve(std::vector<int>(_n + 1, -1)) {
        for (int i = 2; i <= n; i++) {
            if (sieve[i] < 0) {
                sieve[i] = i;
                prime.emplace_back(i);
            }
            for (auto p : prime) {
                if (u64(p) * u64(i) > u64(n) || p > sieve[i]) break;
                sieve[p * i] = p;
            }
        }
    }

    std::vector<int> prime_table() const {
        return prime;
    }

    std::vector<std::pair<int, int>> prime_power_table(int m) const {
        assert(m <= n);
        std::vector<std::pair<int, int>> table(m + 1, {1, 1});
        for (int i = 2; i <= m; i++) {
            int p = sieve[i];
            table[i] = {p, p};
            if (sieve[i / p] == p) {
                table[i] = table[i / p];
                table[i].second *= p;
            }
        }
        return table;
    }

    std::vector<std::pair<int, int>> factorize(int x) {
        assert(x <= n);
        std::vector<std::pair<int, int>> res;
        while (x > 1) {
            int p = sieve[x];
            int exp = 0;
            if (p < 0) {
                res.emplace_back(x, 1);
                break;
            }
            while (sieve[x] == p) {
                x /= p;
                exp++;
            }
            res.emplace_back(p, exp);
        }
        return res;
    }

    std::vector<int> divisors(int x) {
        assert(x <= n);
        std::vector<int> res;
        res.emplace_back(1);
        auto pf = factorize(x);
        for (auto p : pf) {
            int sz = (int)res.size();
            for (int i = 0; i < sz; i++) {
                int ret = 1;
                for (int j = 0; j < p.second; j++) {
                    ret *= p.first;
                    res.emplace_back(res[i] * ret);
                }
            }
        }
        return res;
    }

    template <class T> std::vector<T> fast_zeta(const std::vector<T> &f) {
        std::vector<T> F = f;
        int sz = f.size();
        assert(sz <= n + 1);
        for (int i = 2; i < sz; i++) {
            if (sieve[i] != i) continue;
            for (int j = (sz - 1) / i; j >= 1; j--) {
                F[j] += F[j * i];
            }
        }
        return F;
    }

    template <class T> std::vector<T> fast_mobius(const std::vector<T> &F) {
        std::vector<T> f = F;
        int sz = F.size();
        assert(sz <= n + 1);
        for (int i = 2; i < sz; i++) {
            if (sieve[i] != i) continue;
            for (int j = 1; j * i < sz; j++) {
                f[j] -= f[j * i];
            }
        }
        return f;
    }

    template <Modint mint> std::vector<mint> pow_table(int m, int k) {
        assert(m <= n && k >= 0);
        std::vector<mint> table(m + 1, 1);
        table[0] = (k == 0);
        for (int i = 2; i <= m; i++) {
            if (sieve[i] == i) {
                table[i] = mint(i).pow(k);
                continue;
            }
            table[i] = table[sieve[i]] * table[i / sieve[i]];
        }
        return table;
    }

    template <Modint mint> std::vector<mint> inv_table() {
        return pow_table(mint::mod() - 2);
    }
};

}  // namespace ebi

#line 7 "convolution/dirichlet_convolution.hpp"

namespace ebi {

template <class T>
std::vector<T> dirichlet_convolution(const std::vector<T> &a,
                                     const std::vector<T> &b) {
    assert(a.size() == b.size());
    int n = a.size() - 1;
    std::vector<T> c(n + 1, 0);
    for (int i = 1; i <= n; i++) {
        for (int j = 1; i * j <= n; j++) {
            c[i * j] += a[i] * b[j];
        }
    }
    return c;
}

template <class T>
std::vector<T> dirichlet_convolution_left_is_multiplicative_function(
    const std::vector<T> &a, const std::vector<T> &b) {
    assert(a.size() == b.size());
    int n = a.size() - 1;
    static int m = 1;
    static std::vector<int> primes;
    if (m < n) {
        while (m < n) m <<= 1;
        eratosthenes_sieve sieve(m);
        primes = sieve.prime_table();
    }
    std::vector<T> c = b;
    for (auto p : primes) {
        if (p > n) break;
        for (int i = n / p; i >= 1; i--) {
            int s = p * i;
            int pk = p, j = i;
            while (1) {
                c[s] += a[pk] * c[j];
                if (j % p != 0) break;
                pk *= p;
                j /= p;
            }
        }
    }
    return c;
}

template <class T>
std::vector<T> dirichlet_convolution_multiplicative_function(
    const std::vector<T> &a, const std::vector<T> &b) {
    assert(a.size() == b.size());
    int n = a.size() - 1;
    static int m = 1;
    static std::vector<std::pair<int, int>> prime_pow_table;
    if (m < n) {
        while (m < n) m <<= 1;
        linear_sieve sieve(m);
        prime_pow_table = sieve.prime_power_table(m);
    }
    std::vector<T> c(n + 1, 0);
    c[1] = a[1] * b[1];
    for (int i = 2; i <= n; i++) {
        auto [p, pk] = prime_pow_table[i];
        if (pk == i) {
            for (int j = 1; j <= i; j *= p) {
                c[i] += a[j] * b[i / j];
            }
        } else {
            c[i] = c[i / pk] * c[pk];
        }
    }
    return c;
}

}  // namespace ebi
#line 9 "math/DirichletSeries.hpp"

namespace ebi {

template <class T, int id> struct DirichletSeries {
  private:
    using Self = DirichletSeries<T, id>;

    void set(std::function<T(i64)> f, std::function<T(i64)> F) {
        for (int i = 1; i <= K; i++) {
            a[i] = f(i);
        }
        for (int i = 1; i <= L; i++) {
            A[i] = F(N / i);
        }
    }

  public:
    DirichletSeries() : a(K + 1), A(L + 1) {}

    DirichletSeries(std::function<T(i64)> f, std::function<T(i64)> F,
                    bool _is_multiplicative = false)
        : a(K + 1), A(L + 1), is_multiplicative(_is_multiplicative) {
        set(f, F);
    }

    Self operator+(const Self &rhs) const noexcept {
        return Self(*this) += rhs;
    }
    Self operator-(const Self &rhs) const noexcept {
        return Self(*this) -= rhs;
    }
    Self operator*(const Self &rhs) const noexcept {
        return Self(*this) *= rhs;
    }
    Self operator/(const Self &rhs) const noexcept {
        return Self(*this) /= rhs;
    }

    Self operator+=(const Self &rhs) noexcept {
        for (int i = 1; i <= K; i++) {
            a[i] += rhs.a[i];
        }
        for (int i = 1; i <= L; i++) {
            A[i] += rhs.A[i];
        }
        return *this;
    }

    Self operator-=(const Self &rhs) noexcept {
        for (int i = 1; i <= K; i++) {
            a[i] -= rhs.a[i];
        }
        for (int i = 1; i <= L; i++) {
            A[i] -= rhs.A[i];
        }
        return *this;
    }

    Self operator*=(const Self &rhs) noexcept {
        Self ret;
        if (is_multiplicative && rhs.is_multiplicative) {
            ret.a = dirichlet_convolution_multiplicative_function(a, rhs.a);
            ret.is_multiplicative = true;
        } else if (is_multiplicative) {
            ret.a =
                dirichlet_convolution_left_is_multiplicative_function(a, rhs.a);
        } else if (rhs.is_multiplicative) {
            ret.a =
                dirichlet_convolution_left_is_multiplicative_function(rhs.a, a);
        } else {
            ret.a = dirichlet_convolution(a, rhs.a);
        }
        std::vector<T> sum_a = a, sum_b = rhs.a;
        for (int i = 1; i < K; i++) {
            sum_a[i + 1] += sum_a[i];
            sum_b[i + 1] += sum_b[i];
        }
        auto get_A = [&](i64 x) -> T {
            if (x <= K) {
                return sum_a[x];
            } else {
                return A[N / x];
            }
        };
        auto get_B = [&](i64 x) -> T {
            if (x <= K) {
                return sum_b[x];
            } else {
                return rhs.A[N / x];
            }
        };
        for (i64 l = L, m = 1; l >= 1; l--) {
            i64 n = N / l;
            while (m * m <= n) m++;
            m--;
            for (int i = 1; i <= m; i++) {
                ret.A[l] +=
                    a[i] * get_B(n / i) + (get_A(n / i) - get_A(m)) * rhs.a[i];
            }
        }
        return ret;
    }

    // c = a / b

    Self operator/=(const Self &b) noexcept {
        Self c = *this;
        T inv_a = b.a[1].inv();
        for (int i = 1; i <= K; i++) {
            c.a[i] *= inv_a;
            for (int j = 2; i * j <= K; j++) {
                c.a[i * j] -= c.a[i] * b.a[j];
            }
        }
        std::vector<T> sum_b = b.a, sum_c = c.a;
        for (int i = 1; i < K; ++i) {
            sum_b[i + 1] += sum_b[i];
            sum_c[i + 1] += sum_c[i];
        }
        auto get_B = [&](i64 x) -> T {
            if (x <= K) {
                return sum_b[x];
            } else {
                return b.A[N / x];
            }
        };
        auto get_C = [&](i64 x) -> T {
            if (x <= K) {
                return sum_c[x];
            } else {
                return c.A[N / x];
            }
        };
        for (i64 l = L, m = 1; l >= 1; l--) {
            i64 n = N / l;
            while (m * m <= n) m++;
            m--;
            for (int i = 2; i <= m; i++) {
                c.A[l] -= b.a[i] * get_C(n / i);
            }

            for (int i = 1; i <= m; i++) {
                c.A[l] -= c.a[i] * (get_B(n / i) - get_B(m));
            }
            c.A[l] *= inv_a;
        }
        return *this = c;
    }

    Self pow(u64 n) const {
        Self res;
        res.a[1] = 1;
        res.is_multiplicative = is_multiplicative;
        std::fill(res.A.begin(), res.A.end(), 1);
        Self x = *this;
        while (n > 0) {
            if (n & 1) res = res * x;
            x = x * x;
            n >>= 1;
        }
        return res;
    }

    T get_sum() {
        return A[1];
    }

    static Self zeta() {
        Self ret;
        std::fill(ret.a.begin(), ret.a.end(), 1);
        for (int i = 1; i <= L; i++) {
            ret.A[i] = N / i;
        }
        ret.is_multiplicative = true;
        return ret;
    }

    static Self zeta1() {
        Self ret;
        ret.is_multiplicative = true;
        std::iota(ret.a.begin(), ret.a.end(), 0);
        T inv2 = T(2).inv();
        for (int i = 1; i <= L; i++) {
            i64 n = N / i;
            ret.A[i] = T(n) * T(n + 1) * inv2;
        }
        return ret;
    }

    static Self zeta2() {
        Self ret;
        ret.is_multiplicative = true;
        for (int i = 1; i <= K; i++) {
            ret.a[i] = i * i;
        }
        T inv6 = T(6).inv();
        for (int i = 1; i <= L; i++) {
            i64 n = N / i;
            ret.A[i] = T(n) * T(n + 1) * T(2 * n + 1) * inv6;
        }
    }

    static void set_size(i64 n) {
        N = n;
        if (N <= 10) {
            K = N;
            L = 1;
        } else if (N <= 5000) {
            K = 1;
            while (K * K < N) K++;
            L = (N + K - 1) / K;
        } else {
            L = 1;
            while (L * L * L / 50 < N) L++;
            K = (N + L - 1) / L;
        }
    }

    static void set_size_multiplicative(i64 n) {
        N = n;
        L = 1;
        while (L * L * L < N) L++;
        K = L * L;
    }

  private:
    static i64 N, K, L;
    static std::vector<std::pair<int, int>> prime_pow_table;
    std::vector<T> a, A;
    bool is_multiplicative = false;
};

template <class T, int id> i64 DirichletSeries<T, id>::N = 1000000;
template <class T, int id> i64 DirichletSeries<T, id>::K = 10000;
template <class T, int id> i64 DirichletSeries<T, id>::L = 100;
template <class T, int id>
std::vector<std::pair<int, int>> DirichletSeries<T, id>::prime_pow_table = {};

}  // namespace ebi

Back to top page