segmented_sieve.hpp

SECTIONMath INCLUDEnoya/segmented_sieve.hpp

分段枚举不超过 \(n\) 的所有素数并按序回调；适合 \(n\) 较大但只需区间块内存。

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

Complexity: Time: O(n log log n). Space: O(sqrt(n) + bs).

AC 记录：enumerate_primes。

跳到代码 · GitHub ↗

Implementation

ReferenceSourceExpanded

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

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

#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 val) {
  using u128 = unsigned __int128;
  std::uint64_t rt = std::uint64_t(std::sqrt(static_cast<long double>(val)));
  while (u128(rt + 1) * (rt + 1) <= val) {
    rt++;
  }
  while (u128(rt) * rt > val) {
    rt--;
  }
  return rt;
}

} // 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 &&cal,
                             std::size_t bs = 1U << 20) {
  if (n < 2) {
    return 0;
  }
  bs = std::max<std::size_t>(bs, 1);

  std::uint64_t lim = segmented_sieve_internal::integer_sqrt(n);
  std::vector<unsigned char> cmp(lim + 1);
  std::vector<std::uint64_t> ps;
  for (std::uint64_t val = 2; val <= lim; val++) {
    if (cmp[val]) {
      continue;
    }
    ps.push_back(val);
    if (val <= lim / val) {
      for (std::uint64_t mul = val * val; mul <= lim; mul += val) {
        cmp[mul] = true;
      }
    }
  }

  std::uint64_t np = 0;
  cal(np++, std::uint64_t(2));
  for (std::uint64_t l = 3; l <= n;) {
    std::uint64_t ava = (n - l) / 2 + 1;
    std::size_t cnt =
        std::size_t(std::min<std::uint64_t>(ava, std::uint64_t(bs)));
    std::uint64_t r = l + 2 * (cnt - 1);
    std::vector<unsigned char> com(cnt);

    for (std::uint64_t p : ps) {
      if (p == 2) {
        continue;
      }
      if (p > r / p) {
        break;
      }
      std::uint64_t a = p * p;
      if (a < l) {
        std::uint64_t quo = l / p + (l % p != 0);
        a = quo * p;
      }
      if ((a & 1) == 0) {
        a += p;
      }
      if (a > r) {
        continue;
      }
      std::uint64_t ste = 2 * p;
      for (std::uint64_t mul = a;; mul += ste) {
        com[(mul - l) / 2] = true;
        if (ste > r - mul) {
          break;
        }
      }
    }

    for (std::size_t off = 0; off < cnt; off++) {
      if (!com[off]) {
        cal(np++, l + 2 * off);
      }
    }
    if (r == n || r + 2 < r) {
      break;
    }
    l = r + 2;
  }
  return np;
}

} // namespace noya

#ifndef NOYA_SEGMENTED_SIEVE_HPP
#define NOYA_SEGMENTED_SIEVE_HPP 1

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

#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 val) {
  using u128 = unsigned __int128;
  std::uint64_t rt = std::uint64_t(std::sqrt(static_cast<long double>(val)));
  while (u128(rt + 1) * (rt + 1) <= val) {
    rt++;
  }
  while (u128(rt) * rt > val) {
    rt--;
  }
  return rt;
}

} // 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 &&cal,
                             std::size_t bs = 1U << 20) {
  if (n < 2) {
    return 0;
  }
  bs = std::max<std::size_t>(bs, 1);

  std::uint64_t lim = segmented_sieve_internal::integer_sqrt(n);
  std::vector<unsigned char> cmp(lim + 1);
  std::vector<std::uint64_t> ps;
  for (std::uint64_t val = 2; val <= lim; val++) {
    if (cmp[val]) {
      continue;
    }
    ps.push_back(val);
    if (val <= lim / val) {
      for (std::uint64_t mul = val * val; mul <= lim; mul += val) {
        cmp[mul] = true;
      }
    }
  }

  std::uint64_t np = 0;
  cal(np++, std::uint64_t(2));
  for (std::uint64_t l = 3; l <= n;) {
    std::uint64_t ava = (n - l) / 2 + 1;
    std::size_t cnt =
        std::size_t(std::min<std::uint64_t>(ava, std::uint64_t(bs)));
    std::uint64_t r = l + 2 * (cnt - 1);
    std::vector<unsigned char> com(cnt);

    for (std::uint64_t p : ps) {
      if (p == 2) {
        continue;
      }
      if (p > r / p) {
        break;
      }
      std::uint64_t a = p * p;
      if (a < l) {
        std::uint64_t quo = l / p + (l % p != 0);
        a = quo * p;
      }
      if ((a & 1) == 0) {
        a += p;
      }
      if (a > r) {
        continue;
      }
      std::uint64_t ste = 2 * p;
      for (std::uint64_t mul = a;; mul += ste) {
        com[(mul - l) / 2] = true;
        if (ste > r - mul) {
          break;
        }
      }
    }

    for (std::size_t off = 0; off < cnt; off++) {
      if (!com[off]) {
        cal(np++, l + 2 * off);
      }
    }
    if (r == n || r + 2 < r) {
      break;
    }
    l = r + 2;
  }
  return np;
}

} // 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) + bs).

namespace noya {

namespace segmented_sieve_internal {

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

} // 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 &&cal,
                             std::size_t bs = 1U << 20) {
  if (n < 2) {
    return 0;
  }
  bs = std::max<std::size_t>(bs, 1);

  std::uint64_t lim = segmented_sieve_internal::integer_sqrt(n);
  std::vector<unsigned char> cmp(lim + 1);
  std::vector<std::uint64_t> ps;
  for (std::uint64_t val = 2; val <= lim; val++) {
    if (cmp[val]) {
      continue;
    }
    ps.push_back(val);
    if (val <= lim / val) {
      for (std::uint64_t mul = val * val; mul <= lim; mul += val) {
        cmp[mul] = true;
      }
    }
  }

  std::uint64_t np = 0;
  cal(np++, std::uint64_t(2));
  for (std::uint64_t l = 3; l <= n;) {
    std::uint64_t ava = (n - l) / 2 + 1;
    std::size_t cnt =
        std::size_t(std::min<std::uint64_t>(ava, std::uint64_t(bs)));
    std::uint64_t r = l + 2 * (cnt - 1);
    std::vector<unsigned char> com(cnt);

    for (std::uint64_t p : ps) {
      if (p == 2) {
        continue;
      }
      if (p > r / p) {
        break;
      }
      std::uint64_t a = p * p;
      if (a < l) {
        std::uint64_t quo = l / p + (l % p != 0);
        a = quo * p;
      }
      if ((a & 1) == 0) {
        a += p;
      }
      if (a > r) {
        continue;
      }
      std::uint64_t ste = 2 * p;
      for (std::uint64_t mul = a;; mul += ste) {
        com[(mul - l) / 2] = true;
        if (ste > r - mul) {
          break;
        }
      }
    }

    for (std::size_t off = 0; off < cnt; off++) {
      if (!com[off]) {
        cal(np++, l + 2 * off);
      }
    }
    if (r == n || r + 2 < r) {
      break;
    }
    l = r + 2;
  }
  return np;
}

} // namespace noya