tests/math/enum-primes/libchecker.test.cpp

View this file on GitHub
Last update: 2025-03-03 00:33:48+09:00
Problem: https://judge.yosupo.jp/problem/enumerate_primes

Depends on

素数列挙 (library/math/enum-primes.hpp)

Code

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (n); ++i)
using ll = long long;
using ull = unsigned long long;

#include "../../../cpp/library/math/enum-primes.hpp"
#define PROBLEM "https://judge.yosupo.jp/problem/enumerate_primes"

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  int n, a, b;
  cin >> n >> a >> b;
  vector<int> primes = asalib::math::enum_primes(n);
  vector<int> ans;
  for (int i = b; i < primes.size(); i += a) ans.push_back(primes[i]);
  cout << primes.size() << ' ' << ans.size() << '\n';
  rep(i, ans.size()) cout << ans[i] << " \n"[i + 1 == ans.size()];
  return 0;
}

#line 1 "tests/math/enum-primes/libchecker.test.cpp"
#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (n); ++i)
using ll = long long;
using ull = unsigned long long;

#line 2 "library/math/enum-primes.hpp"

#include <bit>
#include <concepts>
#line 7 "library/math/enum-primes.hpp"
using namespace std;

namespace asalib {
  namespace math {
    template<integral T>
    constexpr vector<T> enum_primes(T n) {
      using u8 = char;
      static_assert(sizeof(u8) == 1, "u8 is not 1 byte");

      // 2, 3, 5 の倍数を除外する -> LCM=30 なので余りの候補: 1, 7, 11, 13, 17, 19, 23, 29
      constexpr array<u8, 8> rem = { 1, 7, 11, 13, 17, 19, 23, 29 };
      // c0[i] = rem[i + 1] - rem[i]
      constexpr array<u8, 8> c0 = { 6, 4, 2, 4, 2, 4, 6, 2 };
      // c1[i][j] = (rem[i] * rem[j + 1] / 30) - (rem[i] * rem[j] / 30)
      constexpr array<array<u8, 8>, 8> c1 = { {
        { 0, 0, 0, 0, 0, 0, 0, 1 },
        { 1, 1, 1, 0, 1, 1, 1, 1 },
        { 2, 2, 0, 2, 0, 2, 2, 1 },
        { 3, 1, 1, 2, 1, 1, 3, 1 },
        { 3, 3, 1, 2, 1, 3, 3, 1 },
        { 4, 2, 2, 2, 2, 2, 4, 1 },
        { 5, 3, 1, 4, 1, 3, 5, 1 },
        { 6, 4, 2, 4, 2, 4, 6, 1 },
      } };
      // rr2idx[i][j] = rem[i] * rem[j] % 30 のインデックス
      constexpr array<array<u8, 8>, 8> rr2idx = { {
        { 0, 1, 2, 3, 4, 5, 6, 7 },
        { 1, 5, 4, 0, 7, 3, 2, 6 },
        { 2, 4, 0, 6, 1, 7, 3, 5 },
        { 3, 0, 6, 5, 2, 1, 7, 4 },
        { 4, 7, 1, 2, 5, 6, 0, 3 },
        { 5, 3, 7, 1, 6, 0, 4, 2 },
        { 6, 2, 3, 7, 0, 4, 5, 1 },
        { 7, 6, 5, 4, 3, 2, 1, 0 },
      } };

      vector<T> primes = { 2, 3, 5 };
      if (n <= 5) [[unlikely]] {
        while (!primes.empty() && primes.back() > n) primes.pop_back();
        return primes;
      }

      const T siz = n / 30 - (n % 30 == 0);
      vector<u8> is_prime(siz + 1, 0b11111111);
      if (n % 30) {
        is_prime.back() = 0;
        rep(i, 8) {
          if (n % 30 >= rem[i]) is_prime.back() |= 1 << i;
        }
      }
      is_prime[0] &= 0b11111110;  // 1 は素数ではない

      for (T i = 0; i <= siz; ++i) {
        for (u8 bit = is_prime[i]; bit;) {
          const u8 idx = countr_zero(unsigned(bit));
          bit ^= 1 << idx;
          const T p = 30 * i + rem[idx];
          primes.push_back(p);
          if (p > n / p) continue;  // p * p > n
          // https://qiita.com/peria/items/54499b9ce9d5c1e93e5a
          T i2 = i * (30 * i + 2 * rem[idx]) + rem[idx] * rem[idx] / 30;
          u8 idx2 = idx;
          for (; i2 <= siz;) {
            is_prime[i2] &= ~(1 << rr2idx[idx][idx2]);
            i2 += i * c0[idx2] + c1[idx][idx2];
            idx2 = (idx2 + 1) & 7;  // = (idx2 + 1) % 8
          }
        }
      }
      return primes;
    }
  }  // namespace math
}  // namespace asalib
#line 8 "tests/math/enum-primes/libchecker.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/enumerate_primes"

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  int n, a, b;
  cin >> n >> a >> b;
  vector<int> primes = asalib::math::enum_primes(n);
  vector<int> ans;
  for (int i = b; i < primes.size(); i += a) ans.push_back(primes[i]);
  cout << primes.size() << ' ' << ans.size() << '\n';
  rep(i, ans.size()) cout << ans[i] << " \n"[i + 1 == ans.size()];
  return 0;
}