素数上の累積和 (src/NumberTheory/sum_on_primes.hpp)

View this file on GitHub
View document part on GitHub
Last update: 2024-02-19 15:31:52+09:00
Include: #include "src/NumberTheory/sum_on_primes.hpp"

$\newcommand{\floor}[1]{\left\lfloor{#1}\right\rfloor}$

加法的関数もしくは乗法的関数 $f$ について，その累積和

$ \displaystyle F(N) = \sum_{i=1}^N f(i) $

を計算するための関数群．

数論的関数 $f$ について，素数上で $n\in \left\lbrace \left. \floor{\frac{N}{x}} \right \vert x\in \mathbb{Z}, 1\leq x\leq N\right\rbrace$ まで累積和をとったもの

$ \displaystyle P(n) = \sum_{\substack{p: \mathrm{prime} \ p \leq n}} f(p) $

を利用することで $F(n)$ を計算する.
$P(n)$ はCumSumQuotient<T> クラスとして扱う.

関数名	概要	計算量
`sums_of_powers_on_primes<T>(N,D)`	$n\in \left\lbrace \left. \floor{\frac{N}{x}} \right \vert x\in \mathbb{Z}, 1\leq x\leq N\right\rbrace$ までの素数上の $k$ 乗数の累積和 $\displaystyle \sum_{\substack{p: \mathrm{prime} \ p \leq n}} p^k$ を $k=0,\dots,D$ まで計算したものを返す. 特に $k=0$ は素数計数関数 $\pi(n)$ となる. 返り値は`vector<CumSumQuotient<T>>`.	$\displaystyle O\left(\frac{DN^{3/4}}{\log N}\right)$
`additive_sum(P,f)`	加法的関数 $f$ について累積和 $F(N)$ を返す. $\displaystyle P(n)=\sum_{\substack{p: \mathrm{prime} \ p \leq n}} f(p)$ を表す `CumSumQuotient<T>` クラスの `P` と $f(p^e)$ を表す `f(p,e)`を渡す．	$\displaystyle O\left(\sqrt{N}\right)$
`multiplicative_sum(P,f)`	乗法的関数 $f$ について累積和 $F(N)$ を返す. $\displaystyle P(n)=\sum_{\substack{p: \mathrm{prime} \ p \leq n}} f(p)$ を表す `CumSumQuotient<T>` クラスの `P` と $f(p^e)$ を表す `f(p,e)`を渡す．	$\displaystyle O\left(\frac{DN^{3/4}}{\log N}\right)$

Verify

Xmas Contest 2019 D - Sum of (-1)^f(n)
Xmas Contest 2019 E - Sum of f(n) (加法的関数)
洛谷 P5325 【模板】Min_25 筛 ($f(p)$ が2次多項式)

参考

https://oi-wiki.org/math/number-theory/min-25/

Depends on

Verified with

Code

#pragma once
#include "src/NumberTheory/enumerate_primes.hpp"
#include "src/NumberTheory/CumSumQuotient.hpp"
template <class T> std::vector<CumSumQuotient<T>> sums_of_powers_on_primes(uint64_t N, size_t D) {
 size_t K= std::sqrt(N);
 std::vector ret(D + 1, CumSumQuotient<T>(N));
 for (size_t n= 1, d= 0; n <= K; ++n, d= 0)
  for (T prd= n; d <= D; prd*= (n + ++d)) ret[d].X[n]= prd / (d + 1);
 for (size_t n= 1, d= 0; n <= K; ++n, d= 0)
  for (T prd= N / n; d <= D; prd*= ((N / n) + ++d)) ret[d].X[n + K]= prd / (d + 1);
 if (D >= 2) {
  std::vector<T> stir(D + 1, 0);
  stir[1]= 1;
  for (size_t d= 2; d <= D; stir[d++]= 1) {
   for (size_t j= d; --j;) stir[j]= stir[j - 1] + stir[j] * (d - 1);
   for (size_t j= 1; j < d; ++j) ret[d].X-= stir[j] * ret[j].X;
  }
 }
 for (size_t d= 0; d <= D; ++d) ret[d].X-= 1;
 for (int p: enumerate_primes(K)) {
  uint64_t q= uint64_t(p) * p, M= N / p;
  T pw= 1;
  for (size_t d= 0, t= K / p, u= std::min<uint64_t>(K, N / q); d <= D; ++d, pw*= p) {
   auto &X= ret[d].X;
   T tk= X[p - 1];
   for (size_t n= 1; n <= t; ++n) X[n + K]-= (X[n * p + K] - tk) * pw;
   for (size_t n= t + 1; n <= u; ++n) X[n + K]-= (X[double(M) / n] - tk) * pw;
   for (uint64_t n= K; n >= q; --n) X[n]-= (X[double(n) / p] - tk) * pw;
  }
 }
 return ret;
}
template <class T, class F> T additive_sum(const CumSumQuotient<T> &P, const F &f) {
 T ret= P.sum();
 for (uint64_t d= 2, nN, nd; nN= double(P.N) / d; d= nd) ret+= P(nN) * ((nd= double(P.N) / nN + 1) - d);
 for (uint64_t p: enumerate_primes(P.K))
  for (uint64_t pw= p * p, e= 2; pw <= P.N; ++e, pw*= p) ret+= (f(p, e) - f(p, e - 1)) * (P.N / pw);
 return ret;
}
template <class T, class F> T multiplicative_sum(CumSumQuotient<T> P, const F &f) {
 auto ps= enumerate_primes(P.K);
 size_t psz= ps.size();
 for (size_t j= psz; j--;) {
  uint64_t p= ps[j], M= P.N / p, q= p * p;
  size_t t= P.K / p, u= std::min<uint64_t>(P.K, P.N / q);
  T tk= P.X[p - 1];
  for (auto i= q; i <= P.K; ++i) P.X[i]+= (P.X[double(i) / p] - tk) * f(p, 1);
  for (size_t i= u; i > t; --i) P.X[i + P.K]+= (P.X[double(M) / i] - tk) * f(p, 1);
  for (size_t i= t; i; --i) P.X[i + P.K]+= (P.X[i * p + P.K] - tk) * f(p, 1);
 }
 P.X+= 1;
 auto dfs= [&](auto &rc, uint64_t n, size_t bg, T cf) -> T {
  if (cf == T(0)) return T(0);
  T ret= cf * P(n);
  for (auto i= bg; i < psz; ++i) {
   uint64_t p= ps[i], q= p * p, nn= n / q;
   if (!nn) break;
   for (int e= 2; nn; nn/= p, ++e) ret+= rc(rc, nn, i + 1, cf * (f(p, e) - f(p, 1) * f(p, e - 1)));
  }
  return ret;
 };
 return dfs(dfs, P.N, 0, 1);
}

#line 2 "src/NumberTheory/enumerate_primes.hpp"
#include <algorithm>
#include <cstdint>
#line 2 "src/Internal/ListRange.hpp"
#include <vector>
#include <iostream>
#include <iterator>
#include <type_traits>
#define _LR(name, IT, CT) \
 template <class T> struct name { \
  using Iterator= typename std::vector<T>::IT; \
  Iterator bg, ed; \
  Iterator begin() const { return bg; } \
  Iterator end() const { return ed; } \
  size_t size() const { return std::distance(bg, ed); } \
  CT &operator[](int i) const { return bg[i]; } \
 }
_LR(ListRange, iterator, T);
_LR(ConstListRange, const_iterator, const T);
#undef _LR
template <class T> struct CSRArray {
 std::vector<T> dat;
 std::vector<int> p;
 size_t size() const { return p.size() - 1; }
 ListRange<T> operator[](int i) { return {dat.begin() + p[i], dat.begin() + p[i + 1]}; }
 ConstListRange<T> operator[](int i) const { return {dat.cbegin() + p[i], dat.cbegin() + p[i + 1]}; }
};
template <template <class> class F, class T> std::enable_if_t<std::disjunction_v<std::is_same<F<T>, ListRange<T>>, std::is_same<F<T>, ConstListRange<T>>, std::is_same<F<T>, CSRArray<T>>>, std::ostream &> operator<<(std::ostream &os, const F<T> &r) {
 os << '[';
 for (int _= 0, __= r.size(); _ < __; ++_) os << (_ ? ", " : "") << r[_];
 return os << ']';
}
#line 5 "src/NumberTheory/enumerate_primes.hpp"
namespace nt_internal {
using namespace std;
vector<int> ps, lf;
void sieve(int N) {
 static int n= 2;
 if (n > N) return;
 if (lf.resize((N + 1) >> 1); n == 2) ps.push_back(n++);
 int M= (N - 1) / 2;
 for (int j= 1, e= ps.size(); j < e; ++j) {
  int p= ps[j];
  if (int64_t(p) * p > N) break;
  for (auto k= int64_t(p) * max(n / p / 2 * 2 + 1, p) / 2; k <= M; k+= p) lf[k]+= p * !lf[k];
 }
 for (; n <= N; n+= 2)
  if (!lf[n >> 1]) {
   ps.push_back(lf[n >> 1]= n);
   for (auto j= int64_t(n) * n / 2; j <= M; j+= n) lf[j]+= n * !lf[j];
  }
}
ConstListRange<int> enumerate_primes() { return {ps.cbegin(), ps.cend()}; }
ConstListRange<int> enumerate_primes(int N) {
 sieve(N);
 return {ps.cbegin(), upper_bound(ps.cbegin(), ps.cend(), N)};
}
int least_prime_factor(int n) { return n & 1 ? sieve(n), lf[(n >> 1)] : 2; }
// f(p,e) := f(p^e)
template <class T, class F> vector<T> completely_multiplicative_table(int N, const F &f) {
 vector<T> ret(N + 1);
 sieve(N);
 for (int n= 3, i= 1; n <= N; n+= 2, ++i) ret[n]= lf[i] == n ? f(n, 1) : ret[lf[i]] * ret[n / lf[i]];
 if (int n= 4; 2 <= N)
  for (T t= ret[2]= f(2, 1); n <= N; n+= 2) ret[n]= t * ret[n >> 1];
 return ret[1]= 1, ret;
}
}
using nt_internal::enumerate_primes, nt_internal::least_prime_factor, nt_internal::completely_multiplicative_table;
// O(N log k / log N + N)
template <class T> static std::vector<T> pow_table(int N, uint64_t k) {
 if (k == 0) return std::vector<T>(N + 1, 1);
 auto f= [k](int p, int) {
  T ret= 1, b= p;
  for (auto e= k;; b*= b) {
   if (e & 1) ret*= b;
   if (!(e>>= 1)) return ret;
  }
 };
 return completely_multiplicative_table<T>(N, f);
}
#line 3 "src/NumberTheory/CumSumQuotient.hpp"
#include <valarray>
template <class T> struct CumSumQuotient {
 uint64_t N;
 size_t K;
 std::valarray<T> X;
 CumSumQuotient(uint64_t N): N(N), K(std::sqrt(N)), X(K + K + 1) {}
 T &operator[](uint64_t i) { return i > K ? X[K + double(N) / i] : X[i]; }
 T operator()(uint64_t i) const { return i > K ? X[K + double(N) / i] : X[i]; }
 CumSumQuotient &operator+=(const CumSumQuotient &r) { return X+= r.X, *this; }
 CumSumQuotient &operator-=(const CumSumQuotient &r) { return X-= r.X, *this; }
 CumSumQuotient &operator*=(T a) { return X*= a, *this; }
 CumSumQuotient operator-() const {
  CumSumQuotient ret= *this;
  return ret.X= -ret.X, ret;
 }
 CumSumQuotient operator+(const CumSumQuotient &r) const { return CumSumQuotient(*this)+= r; }
 CumSumQuotient operator-(const CumSumQuotient &r) const { return CumSumQuotient(*this)-= r; }
 CumSumQuotient operator*(T a) const { return CumSumQuotient(*this)*= a; }
 friend CumSumQuotient operator*(T a, const CumSumQuotient &x) { return x * a; }
 void add(uint64_t i, T v) {
  for (size_t j= std::min<uint64_t>(N / i, K) + K; j >= i; --j) X[j]+= v;
 }
 T sum() const { return X[K + 1]; }
 T sum(uint64_t i) const { return i > K ? X[K + double(N) / i] : X[i]; }
};
#line 4 "src/NumberTheory/sum_on_primes.hpp"
template <class T> std::vector<CumSumQuotient<T>> sums_of_powers_on_primes(uint64_t N, size_t D) {
 size_t K= std::sqrt(N);
 std::vector ret(D + 1, CumSumQuotient<T>(N));
 for (size_t n= 1, d= 0; n <= K; ++n, d= 0)
  for (T prd= n; d <= D; prd*= (n + ++d)) ret[d].X[n]= prd / (d + 1);
 for (size_t n= 1, d= 0; n <= K; ++n, d= 0)
  for (T prd= N / n; d <= D; prd*= ((N / n) + ++d)) ret[d].X[n + K]= prd / (d + 1);
 if (D >= 2) {
  std::vector<T> stir(D + 1, 0);
  stir[1]= 1;
  for (size_t d= 2; d <= D; stir[d++]= 1) {
   for (size_t j= d; --j;) stir[j]= stir[j - 1] + stir[j] * (d - 1);
   for (size_t j= 1; j < d; ++j) ret[d].X-= stir[j] * ret[j].X;
  }
 }
 for (size_t d= 0; d <= D; ++d) ret[d].X-= 1;
 for (int p: enumerate_primes(K)) {
  uint64_t q= uint64_t(p) * p, M= N / p;
  T pw= 1;
  for (size_t d= 0, t= K / p, u= std::min<uint64_t>(K, N / q); d <= D; ++d, pw*= p) {
   auto &X= ret[d].X;
   T tk= X[p - 1];
   for (size_t n= 1; n <= t; ++n) X[n + K]-= (X[n * p + K] - tk) * pw;
   for (size_t n= t + 1; n <= u; ++n) X[n + K]-= (X[double(M) / n] - tk) * pw;
   for (uint64_t n= K; n >= q; --n) X[n]-= (X[double(n) / p] - tk) * pw;
  }
 }
 return ret;
}
template <class T, class F> T additive_sum(const CumSumQuotient<T> &P, const F &f) {
 T ret= P.sum();
 for (uint64_t d= 2, nN, nd; nN= double(P.N) / d; d= nd) ret+= P(nN) * ((nd= double(P.N) / nN + 1) - d);
 for (uint64_t p: enumerate_primes(P.K))
  for (uint64_t pw= p * p, e= 2; pw <= P.N; ++e, pw*= p) ret+= (f(p, e) - f(p, e - 1)) * (P.N / pw);
 return ret;
}
template <class T, class F> T multiplicative_sum(CumSumQuotient<T> P, const F &f) {
 auto ps= enumerate_primes(P.K);
 size_t psz= ps.size();
 for (size_t j= psz; j--;) {
  uint64_t p= ps[j], M= P.N / p, q= p * p;
  size_t t= P.K / p, u= std::min<uint64_t>(P.K, P.N / q);
  T tk= P.X[p - 1];
  for (auto i= q; i <= P.K; ++i) P.X[i]+= (P.X[double(i) / p] - tk) * f(p, 1);
  for (size_t i= u; i > t; --i) P.X[i + P.K]+= (P.X[double(M) / i] - tk) * f(p, 1);
  for (size_t i= t; i; --i) P.X[i + P.K]+= (P.X[i * p + P.K] - tk) * f(p, 1);
 }
 P.X+= 1;
 auto dfs= [&](auto &rc, uint64_t n, size_t bg, T cf) -> T {
  if (cf == T(0)) return T(0);
  T ret= cf * P(n);
  for (auto i= bg; i < psz; ++i) {
   uint64_t p= ps[i], q= p * p, nn= n / q;
   if (!nn) break;
   for (int e= 2; nn; nn/= p, ++e) ret+= rc(rc, nn, i + 1, cf * (f(p, e) - f(p, 1) * f(p, e - 1)));
  }
  return ret;
 };
 return dfs(dfs, P.N, 0, 1);
}