PNSieve
(Number_Theory/PNSieve.hpp)

• View this file on GitHub
• Last update: 2023-02-11 22:28:05+08:00
• Include: #include "Number_Theory/PNSieve.hpp"

PN筛计算积性函数$f$前缀和：$\sum_{i=1}^n f(i)$。

找到积性函数$g$使的$g(p) = f(p)$其中$p$是质数。$G(x)$是 $g$函数的前缀和。

$f,g$以三参数传入，$f(n,p,c) \rightarrow f(p^c) = f(n)$

时间复杂度: $\mathcal{O} (\max(\sqrt{n},n^t (\frac{\zeta(2t)\zeta(3t)}{\zeta(6t)}))$ $n^t$用来计算$G(x)$ ,$\zeta$ 是Riemann zeta function。

Depends on

• Number_Theory/Prime_Sieve.hpp
• Template/Template.hpp

Code

#pragma once

#include "Prime_Sieve.hpp"

template <class T>
T PNSieve(i64 n, const std::function<T(i64)> G,
          const std::function<T(i64, i64, i64)> f,
          const std::function<T(i64, i64, i64)> g) {
    /*
        \sum_{i=1}^n f(i)
        Find g: g(x) = f(x) when x is prime ,and g is multiplicative
        G(x) is the prifix sum of g.
        Time Complexity: O (max(sqrt(n),n^t
       (\zeta(2*t)*\zeta(3*t))/\zeta(6*t))) n^t for call G(x) ,\zeta is Riemann
       zeta function
    */
    const int sq = sqrt(n) + 10;
    auto prime = prime_sieve(sq);
    std::vector<std::vector<T>> h(prime.size(),
                                  std::vector<T>(std::__lg(n) + 2));
    for (int i = 0; i < int(prime.size()); i++) h[i][0] = T(1);
    T ans = 0;

    auto dfs = [&](auto&& self, i64 x, T y, i64 pos) -> void {
        ans = ans + y * G(n / x);
        if (pos >= int(prime.size())) return;
        if (x > n / prime[pos] / prime[pos]) return;
        for (int w = pos; w < int(prime.size()); w++) {
            if (x > n / prime[w] / prime[w]) return;
            i64 bet = x * prime[w];
            for (i64 j = 1; bet <= n; j++, bet *= prime[w]) {
                if (h[w][j] == T()) {
                    T nowf = f(bet / x, prime[w], j);
                    i64 G = prime[w];
                    for (i64 i = 1; i <= j; i++)
                        nowf = nowf - g(G, prime[w], i) * h[w][j - i],
                        G *= prime[w];
                    h[w][j] = nowf;
                }
                if (h[w][j] != T()) {
                    self(self, bet, h[w][j] * y, w + 1);
                }
            }
        }
    };
    dfs(dfs, 1, 1, 0);
    return ans;
}

#line 2 "Number_Theory/PNSieve.hpp"

#line 2 "Number_Theory/Prime_Sieve.hpp"

#line 2 "Template/Template.hpp"

#include <bits/stdc++.h>

using i64 = std::int64_t;
#line 4 "Number_Theory/Prime_Sieve.hpp"

std::vector<int> prime_sieve(int N) {
    std::vector<bool> sieve(N / 3 + 1, 1);
    for (int p = 5, d = 4, i = 1, sqn = sqrt(N); p <= sqn;
         p += d = 6 - d, i++) {
        if (!sieve[i]) continue;
        for (int q = p * p / 3, r = d * p / 3 + (d * p % 3 == 2), s = 2 * p,
                 qe = sieve.size();
             q < qe; q += r = s - r)
            sieve[q] = 0;
    }
    std::vector<int> ret{2, 3};
    for (int p = 5, d = 4, i = 1; p <= N; p += d = 6 - d, i++)
        if (sieve[i]) ret.push_back(p);
    while (!ret.empty() && ret.back() > N) ret.pop_back();
    return ret;
}
#line 4 "Number_Theory/PNSieve.hpp"

template <class T>
T PNSieve(i64 n, const std::function<T(i64)> G,
          const std::function<T(i64, i64, i64)> f,
          const std::function<T(i64, i64, i64)> g) {
    /*
        \sum_{i=1}^n f(i)
        Find g: g(x) = f(x) when x is prime ,and g is multiplicative
        G(x) is the prifix sum of g.
        Time Complexity: O (max(sqrt(n),n^t
       (\zeta(2*t)*\zeta(3*t))/\zeta(6*t))) n^t for call G(x) ,\zeta is Riemann
       zeta function
    */
    const int sq = sqrt(n) + 10;
    auto prime = prime_sieve(sq);
    std::vector<std::vector<T>> h(prime.size(),
                                  std::vector<T>(std::__lg(n) + 2));
    for (int i = 0; i < int(prime.size()); i++) h[i][0] = T(1);
    T ans = 0;

    auto dfs = [&](auto&& self, i64 x, T y, i64 pos) -> void {
        ans = ans + y * G(n / x);
        if (pos >= int(prime.size())) return;
        if (x > n / prime[pos] / prime[pos]) return;
        for (int w = pos; w < int(prime.size()); w++) {
            if (x > n / prime[w] / prime[w]) return;
            i64 bet = x * prime[w];
            for (i64 j = 1; bet <= n; j++, bet *= prime[w]) {
                if (h[w][j] == T()) {
                    T nowf = f(bet / x, prime[w], j);
                    i64 G = prime[w];
                    for (i64 i = 1; i <= j; i++)
                        nowf = nowf - g(G, prime[w], i) * h[w][j - i],
                        G *= prime[w];
                    h[w][j] = nowf;
                }
                if (h[w][j] != T()) {
                    self(self, bet, h[w][j] * y, w + 1);
                }
            }
        }
    };
    dfs(dfs, 1, 1, 0);
    return ans;
}

Back to top page