Tifa's CP Library

:heavy_check_mark: src/test_cpverifier/unit-test/nt/atkin.test.cpp

Depends on

Code

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

#include "../../../code/nt/atkin.hpp"

#include "../base.hpp"

void test() {
  vecu const primes{
      2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211, 223, 227, 229, 233, 239, 241, 251, 257, 263, 269, 271, 277, 281, 283, 293, 307, 311, 313, 317, 331, 337, 347, 349, 353, 359, 367, 373, 379, 383, 389, 397, 401, 409, 419, 421, 431, 433, 439, 443, 449, 457, 461, 463, 467, 479, 487, 491, 499, 503, 509, 521, 523, 541, 547, 557, 563, 569, 571, 577, 587, 593, 599, 601, 607, 613, 617, 619, 631, 641, 643, 647, 653, 659, 661, 673, 677, 683, 691, 701, 709, 719, 727, 733, 739, 743, 751, 757, 761, 769, 773, 787, 797, 809, 811, 821, 823, 827, 829, 839, 853, 857, 859, 863, 877, 881, 883, 887, 907, 911, 919, 929, 937, 941, 947, 953, 967, 971, 977, 983, 991, 997, 1009, 1013, 1019, 1021, 1031, 1033, 1039, 1049, 1051, 1061, 1063, 1069, 1087, 1091, 1093, 1097, 1103, 1109, 1117, 1123, 1129, 1151, 1153, 1163, 1171, 1181, 1187, 1193, 1201, 1213, 1217, 1223, 1229, 1231, 1237, 1249, 1259, 1277, 1279, 1283, 1289, 1291, 1297, 1301, 1303, 1307, 1319, 1321, 1327, 1361, 1367, 1373, 1381, 1399, 1409, 1423, 1427, 1429, 1433, 1439, 1447, 1451, 1453, 1459, 1471, 1481, 1483, 1487, 1489, 1493, 1499, 1511, 1523, 1531, 1543, 1549, 1553, 1559, 1567, 1571, 1579, 1583, 1597, 1601, 1607, 1609, 1613, 1619, 1621, 1627, 1637, 1657, 1663, 1667, 1669, 1693, 1697, 1699, 1709, 1721, 1723, 1733, 1741, 1747, 1753, 1759, 1777, 1783, 1787, 1789, 1801, 1811, 1823, 1831, 1847, 1861, 1867, 1871, 1873, 1877, 1879, 1889, 1901, 1907, 1913, 1931, 1933, 1949, 1951, 1973, 1979, 1987, 1993, 1997, 1999};

  auto got = tifa_libs::math::atkin(2000);
  check(got, primes);
}

int main() {
  auto tcase = tifa_libs::unittest::pre_test();

  switch (tcase) {
    case tifa_libs::unittest::ts_example_00: test(); break;
    case tifa_libs::unittest::ts_random_00: break;
    case tifa_libs::unittest::ts_random_01: break;
    case tifa_libs::unittest::ts_random_02: break;
    case tifa_libs::unittest::ts_random_03: break;
    case tifa_libs::unittest::ts_random_04: break;
    case tifa_libs::unittest::ts_random_05: break;
    case tifa_libs::unittest::ts_random_06: break;
    case tifa_libs::unittest::ts_random_07: break;
    case tifa_libs::unittest::ts_random_08: break;
    case tifa_libs::unittest::ts_random_09: break;
    default: break;
  }

  tifa_libs::unittest::post_test();
}
#line 1 "src/test_cpverifier/unit-test/nt/atkin.test.cpp"
#define UNITTEST
#define PROBLEM "https://judge.yosupo.jp/problem/aplusb"

#line 1 "src/code/nt/atkin.hpp"



#line 1 "src/code/math/isqrt.hpp"



#line 1 "src/code/util/util.hpp"



#include <bits/stdc++.h>

template <class T>
constexpr T abs(T x) { return x < 0 ? -x : x; }

using i8 = int8_t;
using i16 = int16_t;
using i32 = int32_t;
using i64 = int64_t;
using i128 = __int128_t;
using isz = ptrdiff_t;

using u8 = uint8_t;
using u16 = uint16_t;
using u32 = uint32_t;
using u64 = uint64_t;
using u128 = __uint128_t;
using usz = size_t;

using f32 = float;
using f64 = double;
using f128 = long double;

template <class T>
using ptt = std::pair<T, T>;
template <class T>
using pt3 = std::tuple<T, T, T>;
template <class T>
using pt4 = std::tuple<T, T, T, T>;

template <class T, usz N>
using arr = std::array<T, N>;
template <class T>
using vec = std::vector<T>;
template <class T>
using vvec = vec<vec<T>>;
template <class T>
using v3ec = vec<vvec<T>>;
template <class U, class T>
using vecp = vec<std::pair<U, T>>;
template <class U, class T>
using vvecp = vvec<std::pair<U, T>>;
template <class T>
using vecpt = vec<ptt<T>>;
template <class T>
using vvecpt = vvec<ptt<T>>;

template <class T, class C = std::less<T>>
using pq = std::priority_queue<T, vec<T>, C>;
template <class T>
using pqg = std::priority_queue<T, vec<T>, std::greater<T>>;

using strn = std::string;
using strnv = std::string_view;

using vecu = vec<u32>;
using vvecu = vvec<u32>;
using v3ecu = v3ec<u32>;
using vecu64 = vec<u64>;
using vecb = vec<bool>;
using vvecb = vvec<bool>;

#ifdef ONLINE_JUDGE
#undef assert
#define assert(x) 42
#endif

using namespace std::literals;

constexpr i8 operator""_i8(unsigned long long x) { return (i8)x; }
constexpr i16 operator""_i16(unsigned long long x) { return (i16)x; }
constexpr i32 operator""_i32(unsigned long long x) { return (i32)x; }
constexpr i64 operator""_i64(unsigned long long x) { return (i64)x; }
constexpr isz operator""_iz(unsigned long long x) { return (isz)x; }

constexpr u8 operator""_u8(unsigned long long x) { return (u8)x; }
constexpr u16 operator""_u16(unsigned long long x) { return (u16)x; }
constexpr u32 operator""_u32(unsigned long long x) { return (u32)x; }
constexpr u64 operator""_u64(unsigned long long x) { return (u64)x; }
constexpr usz operator""_uz(unsigned long long x) { return (usz)x; }

inline const auto fn_0 = [](auto&&...) {};


#line 5 "src/code/math/isqrt.hpp"

namespace tifa_libs::math {

constexpr u32 isqrt(u64 x) {
  if (!x) return 0;
  int c = i32(std::bit_width(x) - 1) / 2, sh = 31 - c;
  u32 u = [](u64 x) {
    constexpr u8 TAB[192] = {128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 144, 145, 146, 147, 148, 149, 150, 151, 151, 152, 153, 154, 155, 156, 156, 157, 158, 159, 160, 160, 161, 162, 163, 164, 164, 165, 166, 167, 167, 168, 169, 170, 170, 171, 172, 173, 173, 174, 175, 176, 176, 177, 178, 179, 179, 180, 181, 181, 182, 183, 183, 184, 185, 186, 186, 187, 188, 188, 189, 190, 190, 191, 192, 192, 193, 194, 194, 195, 196, 196, 197, 198, 198, 199, 200, 200, 201, 201, 202, 203, 203, 204, 205, 205, 206, 206, 207, 208, 208, 209, 210, 210, 211, 211, 212, 213, 213, 214, 214, 215, 216, 216, 217, 217, 218, 219, 219, 220, 220, 221, 221, 222, 223, 223, 224, 224, 225, 225, 226, 227, 227, 228, 228, 229, 229, 230, 230, 231, 232, 232, 233, 233, 234, 234, 235, 235, 236, 237, 237, 238, 238, 239, 239, 240, 240, 241, 241, 242, 242, 243, 243, 244, 244, 245, 246, 246, 247, 247, 248, 248, 249, 249, 250, 250, 251, 251, 252, 252, 253, 253, 254, 254, 255, 255, 255};
    u32 u = TAB[(x >> 56) - 64];
    u = (u << 7) + (u32)(x >> 41) / u;
    return (u << 15) + (u32)((x >> 17) / u);
  }(x << 2 * sh);
  u >>= sh;
  u -= (u64)u * u > x;
  return u;
}

}  // namespace tifa_libs::math


#line 5 "src/code/nt/atkin.hpp"

namespace tifa_libs::math {

// @return primes in [1, %lim]
constexpr vecu atkin(u32 lim) {
  if (lim < 2) return {};
  if (lim < 3) return {2};
  if (lim < 5) return {2, 3};

  u32 slim = isqrt(lim);
  if (slim * slim < lim) ++slim;
  vecb vis(lim + 1);

  for (u32 i = 0; i < slim; ++i)
    for (u32 j = 0; j < slim; ++j) {
      if (u64 n = 4_u64 * i * i + j * j; n > lim) break;
      else if (u32(n) % 12 == 1 || u32(n) % 12 == 5) vis[n].flip();
    }
  for (u32 i = 0; i < slim; ++i)
    for (u32 j = 0; j < slim; ++j) {
      if (u64 n = 3_u64 * i * i + j * j; n > lim) break;
      else if (u32(n) % 12 == 7) vis[n].flip();
    }
  for (u32 i = 0; i < slim; ++i)
    for (u32 j = i - 1; ~j; --j) {
      if (u64 n = 3_u64 * i * i - j * j; n > lim) break;
      else if (u32(n) % 12 == 11) vis[n].flip();
    }
  for (u32 i = 5; i <= slim; ++i)
    if (vis[i])
      for (u32 k = i * i, j = k; j <= lim; j += k) vis[j] = false;
  vecu ret{2, 3};
  for (u32 x = 5; x <= lim; ++x)
    if (vis[x]) ret.push_back(x);
  return ret;
}

}  // namespace tifa_libs::math


#line 5 "src/test_cpverifier/unit-test/nt/atkin.test.cpp"

#line 1 "src/test_cpverifier/unit-test/base.hpp"



#line 1 "src/code/io/ios128.hpp"



#line 5 "src/code/io/ios128.hpp"

inline std::istream &operator>>(std::istream &is, i128 &n) {
  bool neg = false;
  while (!neg && !isdigit(is.peek())) {
    if (is.peek() == '-') neg = true;
    is.get();
  }
  n = 0;
  while (isdigit(is.peek())) (n *= 10) += is.get() & 15;
  if (neg) n = -n;
  return is;
}
inline std::istream &operator>>(std::istream &is, u128 &n) {
  while (!isdigit(is.peek())) is.get();
  n = 0;
  while (isdigit(is.peek())) (n *= 10) += is.get() & 15;
  return is;
}
inline std::ostream &operator<<(std::ostream &os, u128 n) {
  if (n > 9) os << n / 10;
  os << (uint_fast32_t)(n % 10);
  return os;
}
inline std::ostream &operator<<(std::ostream &os, i128 n) {
  if (n < 0) {
    os << '-';
    n = -n;
  }
  return os << (u128)n;
}


#line 1 "src/code/io/ios_container.hpp"



#line 1 "src/code/util/traits.hpp"



#line 5 "src/code/util/traits.hpp"

namespace tifa_libs {

template <class T>
concept iterable_c = requires(T v) {
  { v.begin() } -> std::same_as<typename T::iterator>;
  { v.end() } -> std::same_as<typename T::iterator>;
};

template <class T>
concept container_c = iterable_c<T> && !std::derived_from<T, std::basic_string<typename T::value_type>>;

template <class T>
constexpr bool is_char_v = std::is_same_v<T, char> || std::is_same_v<T, signed char> || std::is_same_v<T, unsigned char>;
template <class T>
concept char_c = is_char_v<T>;

template <class T>
constexpr bool is_s128_v = std::is_same_v<T, __int128_t> || std::is_same_v<T, __int128>;
template <class T>
concept s128_c = is_s128_v<T>;

template <class T>
constexpr bool is_u128_v = std::is_same_v<T, __uint128_t> || std::is_same_v<T, unsigned __int128>;
template <class T>
concept u128_c = is_u128_v<T>;

template <class T>
constexpr bool is_i128_v = is_s128_v<T> || is_u128_v<T>;
template <class T>
concept i128_c = is_u128_v<T>;

template <class T>
constexpr bool is_int_v = std::is_integral_v<T> || is_i128_v<T>;
template <class T>
concept int_c = is_int_v<T>;

template <class T>
constexpr bool is_sint_v = is_s128_v<T> || (is_int_v<T> && std::is_signed_v<T>);
template <class T>
concept sint_c = is_sint_v<T>;

template <class T>
constexpr bool is_uint_v = is_u128_v<T> || (is_int_v<T> && std::is_unsigned_v<T>);
template <class T>
concept uint_c = is_uint_v<T>;

template <class T>
concept mint_c = requires(T x) {
  { x.mod() } -> uint_c;
  { x.val() } -> uint_c;
};

template <class T>
constexpr bool is_arithm_v = std::is_arithmetic_v<T> || is_int_v<T>;
template <class T>
concept arithm_c = is_arithm_v<T>;

template <class T>
struct to_sint : std::make_signed<T> {};
template <>
struct to_sint<u128> {
  using type = u128;
};
template <>
struct to_sint<i128> {
  using type = u128;
};
template <class T>
using to_sint_t = typename to_sint<T>::type;

template <class T>
struct to_uint : std::make_unsigned<T> {};
template <>
struct to_uint<u128> {
  using type = u128;
};
template <>
struct to_uint<i128> {
  using type = u128;
};
template <class T>
using to_uint_t = typename to_uint<T>::type;

}  // namespace tifa_libs


#line 5 "src/code/io/ios_container.hpp"

template <tifa_libs::container_c T>
std::istream &operator>>(std::istream &is, T &x) {
  for (auto &i : x) is >> i;
  return is;
}
template <tifa_libs::container_c T>
std::ostream &operator<<(std::ostream &os, T const &x) {
  if (x.begin() == x.end()) return os;
  auto it = x.begin();
  os << *it++;
  for (; it != x.end(); ++it) os << ' ' << *it;
  return os;
}


#line 1 "src/code/io/ios_pair.hpp"



#line 5 "src/code/io/ios_pair.hpp"

template <class T, class U>
std::istream &operator>>(std::istream &is, std::pair<T, U> &p) { return is >> p.first >> p.second; }
template <class T, class U>
std::ostream &operator<<(std::ostream &os, std::pair<T, U> const &p) { return os << p.first << ' ' << p.second; }


#line 1 "src/code/io/ios_tuple.hpp"



#line 5 "src/code/io/ios_tuple.hpp"

template <class... Ts>
std::istream &operator>>(std::istream &is, std::tuple<Ts...> &p) {
  std::apply([&](Ts &...targs) { ((is >> targs), ...); }, p);
  return is;
}
template <class... Ts>
std::ostream &operator<<(std::ostream &os, std::tuple<Ts...> const &p) {
  std::apply(
      [&](Ts const &...targs) {
        usz n = 0;
        ((os << targs << (++n != sizeof...(Ts) ? " " : "")), ...);
      },
      p);
  return os;
}


#line 1 "src/code/rand/gen.hpp"



#line 5 "src/code/rand/gen.hpp"

namespace tifa_libs::rand {

template <class Distri>
class Gen {
  std::conditional_t<sizeof(typename Distri::result_type) <= 4, std::mt19937, std::mt19937_64> re;
  Distri dist;

 public:
  using random_engine = decltype(re);
  using distribution = Distri;
  using result_type = typename Distri::result_type;

  constexpr Gen() : re(std::random_device{}()), dist() {}
  constexpr Gen(result_type a, result_type b) : re(std::random_device{}()), dist(a, b) {}

  constexpr void set_range(result_type a, result_type b) { dist = Distri(a, b); }
  constexpr random_engine& rand_eng() { return re; }
  constexpr Distri& distrib() { return dist; }

  void reset_seed() { re.seed((result_type)std::chrono::duration_cast<std::conditional_t<sizeof(typename Distri::result_type) <= 4, std::chrono::seconds, std::chrono::nanoseconds>>(std::chrono::high_resolution_clock::now().time_since_epoch()).count()); }
  constexpr result_type operator()() { return dist(re); }
  result_type next() { return dist(re); }
};

}  // namespace tifa_libs::rand


#line 10 "src/test_cpverifier/unit-test/base.hpp"

namespace tifa_libs::unittest {

namespace detail__ {

template <class T>
strn to_str(T const &x) {
  std::stringstream ss;
  ss << std::fixed << std::setprecision(12) << x;
  return ss.str();
}

template <class T, class... Ts>
void check_(strn const &pretty_func, strn const &got_str, T const &got, strn const &want_str, T const &want, Ts... param) {
  if constexpr (sizeof...(param) == 0) {
    if (got != want) throw std::runtime_error(pretty_func + ": got \"" + got_str + "\" = " + to_str(got) + ", want \"" + want_str + "\" = " + to_str(want));
  } else {
    if (got != want) throw std::runtime_error(pretty_func + ": got \"" + got_str + "\" = " + to_str(got) + ", want \"" + want_str + "\" = " + to_str(want) + " with" + ((" " + param.first + " = " + ::tifa_libs::unittest::detail__::to_str(param.second) + " ;") + ...));
  }
}

template <class... Ts>
void check_bool_(strn const &pretty_func, strn const &expression, bool res, Ts... param) {
  if constexpr (sizeof...(param) == 0) {
    if (!res) throw std::runtime_error(pretty_func + " :\"" + expression + "\" failed");
  } else {
    if (!res) throw std::runtime_error(pretty_func + " :\"" + expression + "\" failed with" + ((" " + param.first + " = " + ::tifa_libs::unittest::detail__::to_str(param.second) + " ;") + ...));
  }
}

}  // namespace detail__

enum TESTCASE { ts_local,
                ts_example_00,
                ts_example_01,
                ts_random_00,
                ts_random_01,
                ts_random_02,
                ts_random_03,
                ts_random_04,
                ts_random_05,
                ts_random_06,
                ts_random_07,
                ts_random_08,
                ts_random_09 };

inline const std::map<ptt<u32>, TESTCASE> testcase_id{
    {{1234, 5678}, ts_example_00},
    {{1000000000, 1000000000}, ts_example_01},
    {{192279220, 156648746}, ts_random_00},
    {{264704197, 120999146}, ts_random_01},
    {{682152023, 451794314}, ts_random_02},
    {{627477696, 504915182}, ts_random_03},
    {{729561619, 415335212}, ts_random_04},
    {{173330281, 220603612}, ts_random_05},
    {{841413509, 58432763}, ts_random_06},
    {{251229786, 256388306}, ts_random_07},
    {{118232767, 222490630}, ts_random_08},
    {{907649120, 290651129}, ts_random_09}};

inline void post_test([[maybe_unused]] ptt<u32> const &p = {0, 0}) {
#ifndef LOCAL_
  static ptt<u32> p_{0, 0};
  if (p.first || p.second) {
    p_ = p;
    return;
  }
  std::cout << p_.first + p_.second << '\n';
  exit(0);
#endif
}

inline TESTCASE pre_test() {
#ifndef LOCAL_
  ptt<u32> p;
  std::cin >> p.first >> p.second;
  post_test(p);
  return testcase_id.at(p);
#else
  return ts_local;
#endif
}

#define check(got, want, ...) ::tifa_libs::unittest::detail__::check_(__PRETTY_FUNCTION__, #got, got, #want, want __VA_OPT__(, ) __VA_ARGS__)
#define check_bool(expression, ...) ::tifa_libs::unittest::detail__::check_bool_(__PRETTY_FUNCTION__, #expression, expression __VA_OPT__(, ) __VA_ARGS__)
#define check_param(x) \
  std::pair<std::string, decltype(x)> { #x " = ", x }

}  // namespace tifa_libs::unittest


#line 7 "src/test_cpverifier/unit-test/nt/atkin.test.cpp"

void test() {
  vecu const primes{
      2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211, 223, 227, 229, 233, 239, 241, 251, 257, 263, 269, 271, 277, 281, 283, 293, 307, 311, 313, 317, 331, 337, 347, 349, 353, 359, 367, 373, 379, 383, 389, 397, 401, 409, 419, 421, 431, 433, 439, 443, 449, 457, 461, 463, 467, 479, 487, 491, 499, 503, 509, 521, 523, 541, 547, 557, 563, 569, 571, 577, 587, 593, 599, 601, 607, 613, 617, 619, 631, 641, 643, 647, 653, 659, 661, 673, 677, 683, 691, 701, 709, 719, 727, 733, 739, 743, 751, 757, 761, 769, 773, 787, 797, 809, 811, 821, 823, 827, 829, 839, 853, 857, 859, 863, 877, 881, 883, 887, 907, 911, 919, 929, 937, 941, 947, 953, 967, 971, 977, 983, 991, 997, 1009, 1013, 1019, 1021, 1031, 1033, 1039, 1049, 1051, 1061, 1063, 1069, 1087, 1091, 1093, 1097, 1103, 1109, 1117, 1123, 1129, 1151, 1153, 1163, 1171, 1181, 1187, 1193, 1201, 1213, 1217, 1223, 1229, 1231, 1237, 1249, 1259, 1277, 1279, 1283, 1289, 1291, 1297, 1301, 1303, 1307, 1319, 1321, 1327, 1361, 1367, 1373, 1381, 1399, 1409, 1423, 1427, 1429, 1433, 1439, 1447, 1451, 1453, 1459, 1471, 1481, 1483, 1487, 1489, 1493, 1499, 1511, 1523, 1531, 1543, 1549, 1553, 1559, 1567, 1571, 1579, 1583, 1597, 1601, 1607, 1609, 1613, 1619, 1621, 1627, 1637, 1657, 1663, 1667, 1669, 1693, 1697, 1699, 1709, 1721, 1723, 1733, 1741, 1747, 1753, 1759, 1777, 1783, 1787, 1789, 1801, 1811, 1823, 1831, 1847, 1861, 1867, 1871, 1873, 1877, 1879, 1889, 1901, 1907, 1913, 1931, 1933, 1949, 1951, 1973, 1979, 1987, 1993, 1997, 1999};

  auto got = tifa_libs::math::atkin(2000);
  check(got, primes);
}

int main() {
  auto tcase = tifa_libs::unittest::pre_test();

  switch (tcase) {
    case tifa_libs::unittest::ts_example_00: test(); break;
    case tifa_libs::unittest::ts_random_00: break;
    case tifa_libs::unittest::ts_random_01: break;
    case tifa_libs::unittest::ts_random_02: break;
    case tifa_libs::unittest::ts_random_03: break;
    case tifa_libs::unittest::ts_random_04: break;
    case tifa_libs::unittest::ts_random_05: break;
    case tifa_libs::unittest::ts_random_06: break;
    case tifa_libs::unittest::ts_random_07: break;
    case tifa_libs::unittest::ts_random_08: break;
    case tifa_libs::unittest::ts_random_09: break;
    default: break;
  }

  tifa_libs::unittest::post_test();
}

Test cases

Env Name Status Elapsed Memory
g++-12 example_00 :heavy_check_mark: AC 10 ms 4 MB
g++-12 example_01 :heavy_check_mark: AC 9 ms 4 MB
g++-12 random_00 :heavy_check_mark: AC 9 ms 4 MB
g++-12 random_01 :heavy_check_mark: AC 9 ms 4 MB
g++-12 random_02 :heavy_check_mark: AC 9 ms 4 MB
g++-12 random_03 :heavy_check_mark: AC 8 ms 4 MB
g++-12 random_04 :heavy_check_mark: AC 8 ms 4 MB
g++-12 random_05 :heavy_check_mark: AC 9 ms 4 MB
g++-12 random_06 :heavy_check_mark: AC 9 ms 4 MB
g++-12 random_07 :heavy_check_mark: AC 10 ms 4 MB
g++-12 random_08 :heavy_check_mark: AC 10 ms 4 MB
g++-12 random_09 :heavy_check_mark: AC 10 ms 4 MB
Back to top page