segmented_sieve.hpp

SECTIONMath INCLUDEnoya/segmented_sieve.hpp

Call callback(index, prime) for every prime not exceeding n and return pi(n). The callback is invoked in increasing prime order.

Verified by enumerate_primes.

\[ \displaystyle \pi(n)=\#\{p\le n: p\text{ prime}\} \]

Implementation

View on GitHub

SourceExpanded

#ifndef NOYA_SEGMENTED_SIEVE_HPP
#define NOYA_SEGMENTED_SIEVE_HPP 1

/// @complexity Time: O(n log log n).
/// Space: O(sqrt(n) + block_size).

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <vector>

namespace noya {

namespace segmented_sieve_internal {

inline std::uint64_t integer_sqrt(std::uint64_t value) {
  using u128 = unsigned __int128;
  std::uint64_t root =
      std::uint64_t(std::sqrt(static_cast<long double>(value)));
  while (u128(root + 1) * (root + 1) <= value) {
    root++;
  }
  while (u128(root) * root > value) {
    root--;
  }
  return root;
}

} // namespace segmented_sieve_internal

/// @brief Call callback(index, prime) for every prime not exceeding n and
/// return pi(n). The callback is invoked in increasing prime order.
template <class Callback>
std::uint64_t for_each_prime(std::uint64_t n, Callback &&callback,
                             std::size_t block_size = 1U << 20) {
  if (n < 2) {
    return 0;
  }
  block_size = std::max<std::size_t>(block_size, 1);

  std::uint64_t limit = segmented_sieve_internal::integer_sqrt(n);
  std::vector<unsigned char> base_composite(limit + 1);
  std::vector<std::uint64_t> base_primes;
  for (std::uint64_t value = 2; value <= limit; value++) {
    if (base_composite[value]) {
      continue;
    }
    base_primes.push_back(value);
    if (value <= limit / value) {
      for (std::uint64_t multiple = value * value; multiple <= limit;
           multiple += value) {
        base_composite[multiple] = true;
      }
    }
  }

  std::uint64_t prime_count = 0;
  callback(prime_count++, std::uint64_t(2));
  for (std::uint64_t left = 3; left <= n;) {
    std::uint64_t available = (n - left) / 2 + 1;
    std::size_t count = std::size_t(
        std::min<std::uint64_t>(available, std::uint64_t(block_size)));
    std::uint64_t right = left + 2 * (count - 1);
    std::vector<unsigned char> composite(count);

    for (std::uint64_t prime : base_primes) {
      if (prime == 2) {
        continue;
      }
      if (prime > right / prime) {
        break;
      }
      std::uint64_t first = prime * prime;
      if (first < left) {
        std::uint64_t quotient = left / prime + (left % prime != 0);
        first = quotient * prime;
      }
      if ((first & 1) == 0) {
        first += prime;
      }
      if (first > right) {
        continue;
      }
      std::uint64_t step = 2 * prime;
      for (std::uint64_t multiple = first;; multiple += step) {
        composite[(multiple - left) / 2] = true;
        if (step > right - multiple) {
          break;
        }
      }
    }

    for (std::size_t offset = 0; offset < count; offset++) {
      if (!composite[offset]) {
        callback(prime_count++, left + 2 * offset);
      }
    }
    if (right == n || right + 2 < right) {
      break;
    }
    left = right + 2;
  }
  return prime_count;
}

} // namespace noya

#endif // NOYA_SEGMENTED_SIEVE_HPP

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <vector>

/// @complexity Time: O(n log log n).
/// Space: O(sqrt(n) + block_size).

namespace noya {

namespace segmented_sieve_internal {

inline std::uint64_t integer_sqrt(std::uint64_t value) {
  using u128 = unsigned __int128;
  std::uint64_t root =
      std::uint64_t(std::sqrt(static_cast<long double>(value)));
  while (u128(root + 1) * (root + 1) <= value) {
    root++;
  }
  while (u128(root) * root > value) {
    root--;
  }
  return root;
}

} // namespace segmented_sieve_internal

/// @brief Call callback(index, prime) for every prime not exceeding n and
/// return pi(n). The callback is invoked in increasing prime order.
template <class Callback>
std::uint64_t for_each_prime(std::uint64_t n, Callback &&callback,
                             std::size_t block_size = 1U << 20) {
  if (n < 2) {
    return 0;
  }
  block_size = std::max<std::size_t>(block_size, 1);

  std::uint64_t limit = segmented_sieve_internal::integer_sqrt(n);
  std::vector<unsigned char> base_composite(limit + 1);
  std::vector<std::uint64_t> base_primes;
  for (std::uint64_t value = 2; value <= limit; value++) {
    if (base_composite[value]) {
      continue;
    }
    base_primes.push_back(value);
    if (value <= limit / value) {
      for (std::uint64_t multiple = value * value; multiple <= limit;
           multiple += value) {
        base_composite[multiple] = true;
      }
    }
  }

  std::uint64_t prime_count = 0;
  callback(prime_count++, std::uint64_t(2));
  for (std::uint64_t left = 3; left <= n;) {
    std::uint64_t available = (n - left) / 2 + 1;
    std::size_t count = std::size_t(
        std::min<std::uint64_t>(available, std::uint64_t(block_size)));
    std::uint64_t right = left + 2 * (count - 1);
    std::vector<unsigned char> composite(count);

    for (std::uint64_t prime : base_primes) {
      if (prime == 2) {
        continue;
      }
      if (prime > right / prime) {
        break;
      }
      std::uint64_t first = prime * prime;
      if (first < left) {
        std::uint64_t quotient = left / prime + (left % prime != 0);
        first = quotient * prime;
      }
      if ((first & 1) == 0) {
        first += prime;
      }
      if (first > right) {
        continue;
      }
      std::uint64_t step = 2 * prime;
      for (std::uint64_t multiple = first;; multiple += step) {
        composite[(multiple - left) / 2] = true;
        if (step > right - multiple) {
          break;
        }
      }
    }

    for (std::size_t offset = 0; offset < count; offset++) {
      if (!composite[offset]) {
        callback(prime_count++, left + 2 * offset);
      }
    }
    if (right == n || right + 2 < right) {
      break;
    }
    left = right + 2;
  }
  return prime_count;
}

} // namespace noya