nt/primesum.hpp

View this file on GitHub
Last update: 2023-08-06 13:04:29+09:00
Include: #include "nt/primesum.hpp"

Depends on

nt/primetable.hpp

Required by

Verified with

Code

#include "nt/primetable.hpp"


/*
N と完全乗法的関数 f の prefix sum 関数 F を与える。
n = floor(N/d) となる n に対する sum_{p <= n} f(p) を計算する。
特に、素数の k 乗和や、mod m ごとでの素数の k 乗和が計算できる。
Complexity: O(N^{3/4}/logN) time, O(N^{1/2}) space.
*/
template <typename T>
struct PrimeSum {
  ll N;
  ll sqN;
  vc<T> sum_lo, sum_hi;
  bool calculated;

  PrimeSum(ll N) : N(N), sqN(sqrtl(N)), calculated(0) {}

  // [1, x] ただし、x = floor(N, i) の形

  T operator[](ll x) {
    assert(calculated);
    return (x <= sqN ? sum_lo[x] : sum_hi[double(N) / x]);
  }

  template <typename F>
  void calc(const F f) {
    auto primes = primetable<int>(sqN);
    sum_lo.resize(sqN + 1);
    sum_hi.resize(sqN + 1);
    FOR3(i, 1, sqN + 1) sum_lo[i] = f(i) - 1;
    FOR3(i, 1, sqN + 1) sum_hi[i] = f(double(N) / i) - 1;
    for (int p: primes) {
      ll pp = ll(p) * p;
      if (pp > N) break;
      int R = min(sqN, N / pp);
      int M = sqN / p;
      T x = sum_lo[p - 1];
      T fp = sum_lo[p] - sum_lo[p - 1];
      for (int i = 1; i <= M; ++i) sum_hi[i] -= fp * (sum_hi[i * p] - x);
      for (int i = M + 1; i <= R; ++i)
        sum_hi[i] -= fp * (sum_lo[N / (double(i) * p)] - x);
      for (int n = sqN; n >= pp; --n) sum_lo[n] -= fp * (sum_lo[n / p] - x);
    }
    calculated = 1;
  }

  void calc_count() {
    calc([](ll x) -> T { return x; });
  }

  void calc_sum() {
    calc([](ll x) -> T {
      ll a = x, b = x + 1;
      if (!(x & 1)) a /= 2;
      if (x & 1) b /= 2;
      return T(a) * T(b);
    });
  }
};

#line 2 "nt/primetable.hpp"

template <typename T = int>
vc<T> primetable(int LIM) {
  ++LIM;
  const int S = 32768;
  static int done = 2;
  static vc<T> primes = {2}, sieve(S + 1);

  if (done < LIM) {
    done = LIM;

    primes = {2}, sieve.assign(S + 1, 0);
    const int R = LIM / 2;
    primes.reserve(int(LIM / log(LIM) * 1.1));
    vc<pair<int, int>> cp;
    for (int i = 3; i <= S; i += 2) {
      if (!sieve[i]) {
        cp.eb(i, i * i / 2);
        for (int j = i * i; j <= S; j += 2 * i) sieve[j] = 1;
      }
    }
    for (int L = 1; L <= R; L += S) {
      array<bool, S> block{};
      for (auto& [p, idx]: cp)
        for (int i = idx; i < S + L; idx = (i += p)) block[i - L] = 1;
      FOR(i, min(S, R - L)) if (!block[i]) primes.eb((L + i) * 2 + 1);
    }
  }
  int k = LB(primes, LIM + 1);
  return {primes.begin(), primes.begin() + k};
}
#line 2 "nt/primesum.hpp"

/*
N と完全乗法的関数 f の prefix sum 関数 F を与える。
n = floor(N/d) となる n に対する sum_{p <= n} f(p) を計算する。
特に、素数の k 乗和や、mod m ごとでの素数の k 乗和が計算できる。
Complexity: O(N^{3/4}/logN) time, O(N^{1/2}) space.
*/
template <typename T>
struct PrimeSum {
  ll N;
  ll sqN;
  vc<T> sum_lo, sum_hi;
  bool calculated;

  PrimeSum(ll N) : N(N), sqN(sqrtl(N)), calculated(0) {}

  // [1, x] ただし、x = floor(N, i) の形

  T operator[](ll x) {
    assert(calculated);
    return (x <= sqN ? sum_lo[x] : sum_hi[double(N) / x]);
  }

  template <typename F>
  void calc(const F f) {
    auto primes = primetable<int>(sqN);
    sum_lo.resize(sqN + 1);
    sum_hi.resize(sqN + 1);
    FOR3(i, 1, sqN + 1) sum_lo[i] = f(i) - 1;
    FOR3(i, 1, sqN + 1) sum_hi[i] = f(double(N) / i) - 1;
    for (int p: primes) {
      ll pp = ll(p) * p;
      if (pp > N) break;
      int R = min(sqN, N / pp);
      int M = sqN / p;
      T x = sum_lo[p - 1];
      T fp = sum_lo[p] - sum_lo[p - 1];
      for (int i = 1; i <= M; ++i) sum_hi[i] -= fp * (sum_hi[i * p] - x);
      for (int i = M + 1; i <= R; ++i)
        sum_hi[i] -= fp * (sum_lo[N / (double(i) * p)] - x);
      for (int n = sqN; n >= pp; --n) sum_lo[n] -= fp * (sum_lo[n / p] - x);
    }
    calculated = 1;
  }

  void calc_count() {
    calc([](ll x) -> T { return x; });
  }

  void calc_sum() {
    calc([](ll x) -> T {
      ll a = x, b = x + 1;
      if (!(x & 1)) a /= 2;
      if (x & 1) b /= 2;
      return T(a) * T(b);
    });
  }
};