Test/LC/enumerate_primes.test.cpp

View this file on GitHub
Last update: 2025-05-09 20:26:39+09:00
Problem: https://judge.yosupo.jp/problem/enumerate_primes

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/enumerate_primes"

#include "../../Src/Number/EnumeratePrimes.hpp"
using namespace zawa;

#include <iostream>

int main() {
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);
    std::ios::sync_with_stdio(false);
    int N, A, B;
    std::cin >> N >> A >> B;
    auto P = zawa::EnumeratePrimes<int>(N);
    std::vector<int> ans;
    ans.reserve(N / A + 1);
    for (int i = B ; i < std::ssize(P) ; i += A) ans.push_back(P[i]);
    std::cout << P.size() << ' ' << ans.size() << '\n';
    for (int i = 0 ; i < std::ssize(ans) ; i++) std::cout << ans[i] << (i + 1 == std::ssize(ans) ? '\n' : ' ');
}

#line 1 "Test/LC/enumerate_primes.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/enumerate_primes"

#line 2 "Src/Number/EnumeratePrimes.hpp"

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 4 "Src/Number/EnumeratePrimes.hpp"

#include <concepts>
#include <vector>
#include <bit>

namespace zawa {

template <std::integral T>
std::vector<T> EnumeratePrimes(usize N) {
    static constexpr u32 p[]{2,3,5,7,11,13,17,19,23,29,31};
    u32 w = 1;
    usize m = 0;
    for ( ; m < std::size(p) and static_cast<u64>(w) * w < N ; w *= p[m++]) ;
    std::vector<u32> R = [&]() {
        std::vector<bool> siv(w, true);
        for (usize i = 0 ; i < m ; i++) {
            for (u32 j = p[i] ; j < w ; j += p[i]) siv[j] = false;
        }
        std::vector<u32> res{1};
        for (u32 i = 2 ; i < w ; i++) if (siv[i]) res.push_back(i);
        return res;
    }();
    std::vector<T> res;
    res.reserve(N / (std::bit_width(N) + 1));
    std::vector<u32> S;
    S.reserve(N / 5u);
    for (u32 kw = 0 ; kw <= N ; kw += w) {
        for (u32 r = 0 ; r < R.size() and kw + R[r] <= N ; r++) {
            S.push_back(kw + R[r]);
        }
    }
    for (usize i = 0 ; i < m ; i++) res.push_back(static_cast<T>(p[i]));
    std::vector<bool> siv(N + 1, true);
    for (usize i = 1 ; i < S.size() ; i++) if (siv[S[i]]) {
        res.push_back(static_cast<T>(S[i])); 
        for (u32 j = i ; j < S.size() ; j++) {
            u64 p = static_cast<u64>(S[i]) * S[j];
            if (p > N) break;
            siv[p] = false;
        }
    }
    return res;
}

} // namespace zawa
#line 4 "Test/LC/enumerate_primes.test.cpp"
using namespace zawa;

#include <iostream>

int main() {
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);
    std::ios::sync_with_stdio(false);
    int N, A, B;
    std::cin >> N >> A >> B;
    auto P = zawa::EnumeratePrimes<int>(N);
    std::vector<int> ans;
    ans.reserve(N / A + 1);
    for (int i = B ; i < std::ssize(P) ; i += A) ans.push_back(P[i]);
    std::cout << P.size() << ' ' << ans.size() << '\n';
    for (int i = 0 ; i < std::ssize(ans) ; i++) std::cout << ans[i] << (i + 1 == std::ssize(ans) ? '\n' : ' ');
}