prime_summatory.hpp

SECTIONMath INCLUDEnoya/prime_summatory.hpp

用商值筛计算 \(\pi(n)\) 及不超过 \(n\) 的素数和，内存仅需约 \(O(\sqrt n)\)。

\[ \displaystyle \pi(n)=\sum_{p\le n}1 \]

Complexity: Time: O(n^(3/4) / log n) arithmetic operations. Space: O(sqrt(n)) quotient blocks.

AC 记录：counting_primes。

跳到代码 · GitHub ↗

Implementation

ReferenceSourceExpanded

当前头文件，省略 include guard；依赖见 #include。

/// @complexity Time: O(n^(3/4) / log n) arithmetic operations.
/// Space: O(sqrt(n)) quotient blocks.

#include <cassert>
#include <cmath>
#include <cstdint>
#include <vector>

namespace noya {

/// @brief Compute pi(n) and the sum of primes at most n with a quotient-block
/// combinatorial sieve using O(sqrt(n)) memory.
struct prime_summatory {
  using u64 = std::uint64_t;
  using u128 = unsigned __int128;

  u64 lim = 0;
  u64 sqr = 0;
  std::vector<u64> vs;
  std::vector<u64> cnt;
  std::vector<u128> sm;
  std::vector<int> ids;
  std::vector<int> idl;

  prime_summatory() = default;
  explicit prime_summatory(u64 n) { build(n); }

  void build(u64 n) {
    lim = n;
    sqr = integer_sqrt(n);
    vs.clear();
    for (u64 l = 1, r; l <= n; l = r + 1) {
      u64 quo = n / l;
      vs.push_back(quo);
      r = n / quo;
    }
    cnt.resize(vs.size());
    sm.resize(vs.size());
    ids.assign(sqr + 1, -1);
    idl.assign(sqr + 1, -1);
    for (int idx = 0; idx < int(vs.size()); idx++) {
      u64 val = vs[idx];
      cnt[idx] = val >= 1 ? val - 1 : 0;
      sm[idx] = val >= 1 ? u128(val) * (u128(val) + 1) / 2 - 1 : 0;
      if (val <= sqr) {
        ids[val] = idx;
      } else {
        idl[n / val] = idx;
      }
    }
    for (u64 p = 2; p <= sqr; p++) {
      int pin = index_of(p);
      int pi = index_of(p - 1);
      if (cnt[pin] == cnt[pi]) {
        continue;
      }
      u64 pc = cnt[pi];
      u128 ps = sm[pi];
      u128 p2 = u128(p) * p;
      for (int idx = 0; idx < int(vs.size()) && u128(vs[idx]) >= p2; idx++) {
        int ri = index_of(vs[idx] / p);
        cnt[idx] -= cnt[ri] - pc;
        sm[idx] -= u128(p) * (sm[ri] - ps);
      }
    }
  }

  u64 count() const { return cnt.empty() ? 0 : cnt.front(); }
  u128 sum() const { return sm.empty() ? 0 : sm.front(); }

private:
  static u64 integer_sqrt(u64 val) {
    u64 rt = u64(std::sqrt(static_cast<long double>(val)));
    while (u128(rt + 1) * (rt + 1) <= val) {
      rt++;
    }
    while (u128(rt) * rt > val) {
      rt--;
    }
    return rt;
  }

  int index_of(u64 val) const {
    assert(val <= lim);
    int idx = val <= sqr ? ids[val] : idl[lim / val];
    assert(idx >= 0 && vs[idx] == val);
    return idx;
  }
};

inline std::uint64_t prime_count(std::uint64_t n) {
  return prime_summatory(n).count();
}

inline unsigned __int128 prime_sum(std::uint64_t n) {
  return prime_summatory(n).sum();
}

} // namespace noya

#ifndef NOYA_PRIME_SUMMATORY_HPP
#define NOYA_PRIME_SUMMATORY_HPP 1

/// @complexity Time: O(n^(3/4) / log n) arithmetic operations.
/// Space: O(sqrt(n)) quotient blocks.

#include <cassert>
#include <cmath>
#include <cstdint>
#include <vector>

namespace noya {

/// @brief Compute pi(n) and the sum of primes at most n with a quotient-block
/// combinatorial sieve using O(sqrt(n)) memory.
struct prime_summatory {
  using u64 = std::uint64_t;
  using u128 = unsigned __int128;

  u64 lim = 0;
  u64 sqr = 0;
  std::vector<u64> vs;
  std::vector<u64> cnt;
  std::vector<u128> sm;
  std::vector<int> ids;
  std::vector<int> idl;

  prime_summatory() = default;
  explicit prime_summatory(u64 n) { build(n); }

  void build(u64 n) {
    lim = n;
    sqr = integer_sqrt(n);
    vs.clear();
    for (u64 l = 1, r; l <= n; l = r + 1) {
      u64 quo = n / l;
      vs.push_back(quo);
      r = n / quo;
    }
    cnt.resize(vs.size());
    sm.resize(vs.size());
    ids.assign(sqr + 1, -1);
    idl.assign(sqr + 1, -1);
    for (int idx = 0; idx < int(vs.size()); idx++) {
      u64 val = vs[idx];
      cnt[idx] = val >= 1 ? val - 1 : 0;
      sm[idx] = val >= 1 ? u128(val) * (u128(val) + 1) / 2 - 1 : 0;
      if (val <= sqr) {
        ids[val] = idx;
      } else {
        idl[n / val] = idx;
      }
    }
    for (u64 p = 2; p <= sqr; p++) {
      int pin = index_of(p);
      int pi = index_of(p - 1);
      if (cnt[pin] == cnt[pi]) {
        continue;
      }
      u64 pc = cnt[pi];
      u128 ps = sm[pi];
      u128 p2 = u128(p) * p;
      for (int idx = 0; idx < int(vs.size()) && u128(vs[idx]) >= p2; idx++) {
        int ri = index_of(vs[idx] / p);
        cnt[idx] -= cnt[ri] - pc;
        sm[idx] -= u128(p) * (sm[ri] - ps);
      }
    }
  }

  u64 count() const { return cnt.empty() ? 0 : cnt.front(); }
  u128 sum() const { return sm.empty() ? 0 : sm.front(); }

private:
  static u64 integer_sqrt(u64 val) {
    u64 rt = u64(std::sqrt(static_cast<long double>(val)));
    while (u128(rt + 1) * (rt + 1) <= val) {
      rt++;
    }
    while (u128(rt) * rt > val) {
      rt--;
    }
    return rt;
  }

  int index_of(u64 val) const {
    assert(val <= lim);
    int idx = val <= sqr ? ids[val] : idl[lim / val];
    assert(idx >= 0 && vs[idx] == val);
    return idx;
  }
};

inline std::uint64_t prime_count(std::uint64_t n) {
  return prime_summatory(n).count();
}

inline unsigned __int128 prime_sum(std::uint64_t n) {
  return prime_summatory(n).sum();
}

} // namespace noya

#endif // NOYA_PRIME_SUMMATORY_HPP

#include <cassert>
#include <cmath>
#include <cstdint>
#include <vector>

/// @complexity Time: O(n^(3/4) / log n) arithmetic operations.
/// Space: O(sqrt(n)) quotient blocks.

namespace noya {

/// @brief Compute pi(n) and the sum of primes at most n with a quotient-block
/// combinatorial sieve using O(sqrt(n)) memory.
struct prime_summatory {
  using u64 = std::uint64_t;
  using u128 = unsigned __int128;

  u64 lim = 0;
  u64 sqr = 0;
  std::vector<u64> vs;
  std::vector<u64> cnt;
  std::vector<u128> sm;
  std::vector<int> ids;
  std::vector<int> idl;

  prime_summatory() = default;
  explicit prime_summatory(u64 n) { build(n); }

  void build(u64 n) {
    lim = n;
    sqr = integer_sqrt(n);
    vs.clear();
    for (u64 l = 1, r; l <= n; l = r + 1) {
      u64 quo = n / l;
      vs.push_back(quo);
      r = n / quo;
    }
    cnt.resize(vs.size());
    sm.resize(vs.size());
    ids.assign(sqr + 1, -1);
    idl.assign(sqr + 1, -1);
    for (int idx = 0; idx < int(vs.size()); idx++) {
      u64 val = vs[idx];
      cnt[idx] = val >= 1 ? val - 1 : 0;
      sm[idx] = val >= 1 ? u128(val) * (u128(val) + 1) / 2 - 1 : 0;
      if (val <= sqr) {
        ids[val] = idx;
      } else {
        idl[n / val] = idx;
      }
    }
    for (u64 p = 2; p <= sqr; p++) {
      int pin = index_of(p);
      int pi = index_of(p - 1);
      if (cnt[pin] == cnt[pi]) {
        continue;
      }
      u64 pc = cnt[pi];
      u128 ps = sm[pi];
      u128 p2 = u128(p) * p;
      for (int idx = 0; idx < int(vs.size()) && u128(vs[idx]) >= p2; idx++) {
        int ri = index_of(vs[idx] / p);
        cnt[idx] -= cnt[ri] - pc;
        sm[idx] -= u128(p) * (sm[ri] - ps);
      }
    }
  }

  u64 count() const { return cnt.empty() ? 0 : cnt.front(); }
  u128 sum() const { return sm.empty() ? 0 : sm.front(); }

private:
  static u64 integer_sqrt(u64 val) {
    u64 rt = u64(std::sqrt(static_cast<long double>(val)));
    while (u128(rt + 1) * (rt + 1) <= val) {
      rt++;
    }
    while (u128(rt) * rt > val) {
      rt--;
    }
    return rt;
  }

  int index_of(u64 val) const {
    assert(val <= lim);
    int idx = val <= sqr ? ids[val] : idl[lim / val];
    assert(idx >= 0 && vs[idx] == val);
    return idx;
  }
};

inline std::uint64_t prime_count(std::uint64_t n) {
  return prime_summatory(n).count();
}

inline unsigned __int128 prime_sum(std::uint64_t n) {
  return prime_summatory(n).sum();
}

} // namespace noya