sum_two_squares.hpp¶
Decide and construct n = a^2 + b^2 for an unsigned 64-bit integer; also count ordered signed representations using the two-square theorem.
Verified by two_square_sum.
\[
\displaystyle n=a^2+b^2
\]
Implementation¶
#ifndef NOYA_SUM_TWO_SQUARES_HPP
#define NOYA_SUM_TWO_SQUARES_HPP 1
/// @complexity Time: Expected integer-factorization time plus O(r log n) to
/// enumerate r representations. Space: O(r + log n).
#include "noya/factorize.hpp"
#include "noya/mod_sqrt.hpp"
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <optional>
#include <utility>
#include <vector>
namespace noya {
/// @brief Decide and construct n = a^2 + b^2 for an unsigned 64-bit integer;
/// also count ordered signed representations using the two-square theorem.
namespace sum_two_squares_internal {
using u64 = std::uint64_t;
using u128 = unsigned __int128;
using i128 = __int128;
inline u64 integer_sqrt(u64 value) {
u64 root = u64(std::sqrt(static_cast<long double>(value)));
while (u128(root + 1) * (root + 1) <= value) {
root++;
}
while (u128(root) * root > value) {
root--;
}
return root;
}
inline std::pair<u64, u64> multiply(std::pair<u64, u64> first,
std::pair<u64, u64> second) {
i128 real =
i128(first.first) * second.first - i128(first.second) * second.second;
i128 imaginary =
i128(first.first) * second.second + i128(first.second) * second.first;
u64 a = u64(real < 0 ? -real : real);
u64 b = u64(imaginary < 0 ? -imaginary : imaginary);
if (a > b) {
std::swap(a, b);
}
return {a, b};
}
inline std::optional<std::pair<u64, u64>> prime_representation(u64 prime) {
assert(prime % 4 == 1 && is_prime(prime));
auto square_root = mod_sqrt(prime - 1, prime);
assert(square_root.has_value());
for (u64 root : {*square_root, prime - *square_root}) {
u64 previous = prime;
u64 current = root;
while (u128(current) * current > prime) {
u64 next = previous % current;
previous = current;
current = next;
}
u64 other_square = prime - current * current;
u64 other = integer_sqrt(other_square);
if (other * other == other_square) {
return std::pair<u64, u64>{std::min(current, other),
std::max(current, other)};
}
}
return std::nullopt;
}
inline std::pair<u64, u64> power(std::pair<u64, u64> base, int exponent) {
std::pair<u64, u64> result{1, 0};
while (exponent > 0) {
if (exponent & 1) {
result = multiply(result, base);
}
exponent >>= 1;
if (exponent > 0) {
base = multiply(base, base);
}
}
return result;
}
using gaussian = std::pair<i128, i128>;
inline gaussian signed_multiply(gaussian first, gaussian second) {
return {first.first * second.first - first.second * second.second,
first.first * second.second + first.second * second.first};
}
inline gaussian signed_power(gaussian base, int exponent) {
gaussian result{1, 0};
while (exponent > 0) {
if (exponent & 1) {
result = signed_multiply(result, base);
}
exponent >>= 1;
if (exponent > 0) {
base = signed_multiply(base, base);
}
}
return result;
}
} // namespace sum_two_squares_internal
inline bool is_sum_two_squares(std::uint64_t n) {
if (n == 0) {
return true;
}
for (auto [prime, exponent] : factorize(n)) {
if (prime % 4 == 3 && exponent % 2 == 1) {
return false;
}
}
return true;
}
/// @brief Return one pair 0 <= a <= b with a^2 + b^2 = n, or nullopt.
inline std::optional<std::pair<std::uint64_t, std::uint64_t>>
sum_two_squares(std::uint64_t n) {
using namespace sum_two_squares_internal;
if (n == 0) {
return std::pair<u64, u64>{0, 0};
}
std::pair<u64, u64> result{1, 0};
for (auto [prime, exponent] : factorize(n)) {
if (prime == 2) {
result = multiply(result, power({1, 1}, exponent));
} else if (prime % 4 == 1) {
auto representation = prime_representation(prime);
assert(representation.has_value());
result = multiply(result, power(*representation, exponent));
} else {
if (exponent % 2 == 1) {
return std::nullopt;
}
u64 scale = 1;
for (int i = 0; i < exponent / 2; i++) {
scale *= prime;
}
result = multiply(result, {scale, 0});
}
}
if (result.first > result.second) {
std::swap(result.first, result.second);
}
return result;
}
/// @brief Count integer pairs (a,b), including signs and order, satisfying
/// a^2 + b^2 = n.
inline unsigned __int128 sum_two_squares_representation_count(std::uint64_t n) {
if (n == 0) {
return 1;
}
unsigned __int128 result = 4;
for (auto [prime, exponent] : factorize(n)) {
if (prime % 4 == 3 && exponent % 2 == 1) {
return 0;
}
if (prime % 4 == 1) {
result *= exponent + 1;
}
}
return result;
}
/// @brief Return every ordered non-negative pair (a,b) with a^2 + b^2 = n.
/// The result is sorted and contains no duplicates. Requires n <= 1e18.
inline std::vector<std::pair<std::uint64_t, std::uint64_t>>
all_sum_two_squares(std::uint64_t n) {
using namespace sum_two_squares_internal;
assert(n <= 1000000000000000000ULL);
if (n == 0) {
return {{0, 0}};
}
auto factors = factorize(n);
for (auto [prime, exponent] : factors) {
if (prime % 4 == 3 && exponent % 2 == 1) {
return {};
}
}
std::vector<gaussian> representatives{{1, 0}};
for (auto [prime, exponent] : factors) {
if (prime % 4 == 3) {
i128 scale = 1;
for (int i = 0; i < exponent / 2; i++) {
scale *= prime;
}
for (auto &[real, imaginary] : representatives) {
real *= scale;
imaginary *= scale;
}
continue;
}
gaussian prime_factor;
if (prime == 2) {
prime_factor = {1, 1};
gaussian multiplier = signed_power(prime_factor, exponent);
for (auto &value : representatives) {
value = signed_multiply(value, multiplier);
}
continue;
}
auto representation = prime_representation(prime);
assert(representation.has_value());
prime_factor = {representation->first, representation->second};
std::vector<gaussian> powers(exponent + 1, {1, 0});
for (int i = 0; i < exponent; i++) {
powers[i + 1] = signed_multiply(powers[i], prime_factor);
}
std::vector<gaussian> next;
next.reserve(representatives.size() * (exponent + 1));
for (gaussian current : representatives) {
for (int chosen = 0; chosen <= exponent; chosen++) {
gaussian conjugate = powers[exponent - chosen];
conjugate.second = -conjugate.second;
gaussian factor = signed_multiply(powers[chosen], conjugate);
next.push_back(signed_multiply(current, factor));
}
}
representatives.swap(next);
}
std::vector<std::pair<u64, u64>> result;
for (auto [real, imaginary] : representatives) {
while (real <= 0 || imaginary < 0) {
i128 old_real = real;
real = -imaginary;
imaginary = old_real;
}
result.emplace_back(u64(real), u64(imaginary));
if (imaginary == 0) {
result.emplace_back(0, u64(real));
}
}
std::sort(result.begin(), result.end());
result.erase(std::unique(result.begin(), result.end()), result.end());
return result;
}
} // namespace noya
#endif // NOYA_SUM_TWO_SQUARES_HPP
#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <numeric>
#include <optional>
#include <utility>
#include <vector>
/// @complexity Time: Expected integer-factorization time plus O(r log n) to
/// enumerate r representations. Space: O(r + log n).
/// @complexity Time: O(log^3 n) primality testing; Pollard-rho factorization is expected about O(n^(1/4)).
/// Space: O(log n) recursion and factors.
namespace noya {
namespace factorize_internal {
using u64 = std::uint64_t;
using u128 = unsigned __int128;
inline u64 multiply_mod(u64 a, u64 b, u64 mod) {
return u64(u128(a) * b % mod);
}
inline u64 power_mod(u64 a, u64 exponent, u64 mod) {
u64 result = 1;
while (exponent > 0) {
if (exponent & 1) {
result = multiply_mod(result, a, mod);
}
a = multiply_mod(a, a, mod);
exponent >>= 1;
}
return result;
}
inline bool miller_rabin(u64 n) {
if (n < 2) {
return false;
}
for (u64 p :
std::array<u64, 12>{2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37}) {
if (n % p == 0) {
return n == p;
}
}
int shift = __builtin_ctzll(n - 1);
u64 odd = (n - 1) >> shift;
for (u64 base :
std::array<u64, 7>{2, 325, 9375, 28178, 450775, 9780504, 1795265022}) {
if (base % n == 0) {
continue;
}
u64 value = power_mod(base % n, odd, n);
if (value == 1 || value == n - 1) {
continue;
}
bool composite = true;
for (int i = 1; i < shift; i++) {
value = multiply_mod(value, value, n);
if (value == n - 1) {
composite = false;
break;
}
}
if (composite) {
return false;
}
}
return true;
}
inline u64 splitmix64(u64 &state) {
u64 z = (state += 0x9e3779b97f4a7c15ULL);
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
return z ^ (z >> 31);
}
inline u64 pollard_rho(u64 n) {
if (n % 2 == 0) {
return 2;
}
if (n % 3 == 0) {
return 3;
}
static u64 state = 0x123456789abcdef0ULL;
while (true) {
u64 y = splitmix64(state) % (n - 1) + 1;
u64 c = splitmix64(state) % (n - 1) + 1;
constexpr u64 block = 128;
u64 g = 1;
u64 r = 1;
u64 q = 1;
u64 x = 0;
u64 saved_y = 0;
auto next = [&](u64 value) {
return u64((u128(multiply_mod(value, value, n)) + c) % n);
};
while (g == 1) {
x = y;
for (u64 i = 0; i < r; i++) {
y = next(y);
}
for (u64 offset = 0; offset < r && g == 1; offset += block) {
saved_y = y;
for (u64 i = 0; i < std::min(block, r - offset); i++) {
y = next(y);
u64 difference = x > y ? x - y : y - x;
q = multiply_mod(q, difference, n);
}
g = std::gcd(q, n);
}
r <<= 1;
}
if (g == n) {
do {
saved_y = next(saved_y);
u64 difference = x > saved_y ? x - saved_y : saved_y - x;
g = std::gcd(difference, n);
} while (g == 1);
}
if (g != n) {
return g;
}
}
}
inline void collect_factors(u64 n, std::vector<u64> &result) {
if (n == 1) {
return;
}
if (miller_rabin(n)) {
result.push_back(n);
return;
}
u64 factor = pollard_rho(n);
collect_factors(factor, result);
collect_factors(n / factor, result);
}
} // namespace factorize_internal
/// @brief Deterministic Miller-Rabin primality test for unsigned 64-bit
/// integers.
inline bool is_prime(std::uint64_t n) {
return factorize_internal::miller_rabin(n);
}
/// @brief Return the prime factors of n with multiplicity in increasing order.
inline std::vector<std::uint64_t> prime_factors(std::uint64_t n) {
assert(n >= 1);
std::vector<std::uint64_t> result;
factorize_internal::collect_factors(n, result);
std::sort(result.begin(), result.end());
return result;
}
/// @brief Return the prime factorization of n as (prime, exponent) pairs.
inline std::vector<std::pair<std::uint64_t, int>> factorize(std::uint64_t n) {
std::vector<std::pair<std::uint64_t, int>> result;
for (std::uint64_t p : prime_factors(n)) {
if (result.empty() || result.back().first != p) {
result.emplace_back(p, 1);
} else {
result.back().second++;
}
}
return result;
}
} // namespace noya
/// @complexity Time: O(log^2 p).
/// Space: O(1).
namespace noya {
/// @brief Compute the smaller square root modulo a prime, or nullopt if no
/// square root exists.
inline std::optional<std::uint64_t> mod_sqrt(std::uint64_t value,
std::uint64_t modulus) {
assert(modulus >= 2 && is_prime(modulus));
value %= modulus;
if (modulus == 2 || value == 0) {
return value;
}
using factorize_internal::multiply_mod;
using factorize_internal::power_mod;
if (power_mod(value, (modulus - 1) / 2, modulus) != 1) {
return std::nullopt;
}
if (modulus % 4 == 3) {
std::uint64_t root = power_mod(value, (modulus + 1) / 4, modulus);
return std::min(root, modulus - root);
}
std::uint64_t odd = modulus - 1;
int exponent = 0;
while ((odd & 1) == 0) {
odd >>= 1;
exponent++;
}
std::uint64_t non_residue = 2;
while (power_mod(non_residue, (modulus - 1) / 2, modulus) != modulus - 1) {
non_residue++;
}
std::uint64_t root = power_mod(value, (odd + 1) / 2, modulus);
std::uint64_t remainder = power_mod(value, odd, modulus);
std::uint64_t step = power_mod(non_residue, odd, modulus);
int remaining = exponent;
while (remainder != 1) {
std::uint64_t squared = remainder;
int shift = 0;
while (squared != 1 && shift < remaining) {
squared = multiply_mod(squared, squared, modulus);
shift++;
}
assert(shift < remaining);
std::uint64_t multiplier =
power_mod(step, std::uint64_t(1) << (remaining - shift - 1), modulus);
root = multiply_mod(root, multiplier, modulus);
step = multiply_mod(multiplier, multiplier, modulus);
remainder = multiply_mod(remainder, step, modulus);
remaining = shift;
}
return std::min(root, modulus - root);
}
} // namespace noya
namespace noya {
/// @brief Decide and construct n = a^2 + b^2 for an unsigned 64-bit integer;
/// also count ordered signed representations using the two-square theorem.
namespace sum_two_squares_internal {
using u64 = std::uint64_t;
using u128 = unsigned __int128;
using i128 = __int128;
inline u64 integer_sqrt(u64 value) {
u64 root = u64(std::sqrt(static_cast<long double>(value)));
while (u128(root + 1) * (root + 1) <= value) {
root++;
}
while (u128(root) * root > value) {
root--;
}
return root;
}
inline std::pair<u64, u64> multiply(std::pair<u64, u64> first,
std::pair<u64, u64> second) {
i128 real =
i128(first.first) * second.first - i128(first.second) * second.second;
i128 imaginary =
i128(first.first) * second.second + i128(first.second) * second.first;
u64 a = u64(real < 0 ? -real : real);
u64 b = u64(imaginary < 0 ? -imaginary : imaginary);
if (a > b) {
std::swap(a, b);
}
return {a, b};
}
inline std::optional<std::pair<u64, u64>> prime_representation(u64 prime) {
assert(prime % 4 == 1 && is_prime(prime));
auto square_root = mod_sqrt(prime - 1, prime);
assert(square_root.has_value());
for (u64 root : {*square_root, prime - *square_root}) {
u64 previous = prime;
u64 current = root;
while (u128(current) * current > prime) {
u64 next = previous % current;
previous = current;
current = next;
}
u64 other_square = prime - current * current;
u64 other = integer_sqrt(other_square);
if (other * other == other_square) {
return std::pair<u64, u64>{std::min(current, other),
std::max(current, other)};
}
}
return std::nullopt;
}
inline std::pair<u64, u64> power(std::pair<u64, u64> base, int exponent) {
std::pair<u64, u64> result{1, 0};
while (exponent > 0) {
if (exponent & 1) {
result = multiply(result, base);
}
exponent >>= 1;
if (exponent > 0) {
base = multiply(base, base);
}
}
return result;
}
using gaussian = std::pair<i128, i128>;
inline gaussian signed_multiply(gaussian first, gaussian second) {
return {first.first * second.first - first.second * second.second,
first.first * second.second + first.second * second.first};
}
inline gaussian signed_power(gaussian base, int exponent) {
gaussian result{1, 0};
while (exponent > 0) {
if (exponent & 1) {
result = signed_multiply(result, base);
}
exponent >>= 1;
if (exponent > 0) {
base = signed_multiply(base, base);
}
}
return result;
}
} // namespace sum_two_squares_internal
inline bool is_sum_two_squares(std::uint64_t n) {
if (n == 0) {
return true;
}
for (auto [prime, exponent] : factorize(n)) {
if (prime % 4 == 3 && exponent % 2 == 1) {
return false;
}
}
return true;
}
/// @brief Return one pair 0 <= a <= b with a^2 + b^2 = n, or nullopt.
inline std::optional<std::pair<std::uint64_t, std::uint64_t>>
sum_two_squares(std::uint64_t n) {
using namespace sum_two_squares_internal;
if (n == 0) {
return std::pair<u64, u64>{0, 0};
}
std::pair<u64, u64> result{1, 0};
for (auto [prime, exponent] : factorize(n)) {
if (prime == 2) {
result = multiply(result, power({1, 1}, exponent));
} else if (prime % 4 == 1) {
auto representation = prime_representation(prime);
assert(representation.has_value());
result = multiply(result, power(*representation, exponent));
} else {
if (exponent % 2 == 1) {
return std::nullopt;
}
u64 scale = 1;
for (int i = 0; i < exponent / 2; i++) {
scale *= prime;
}
result = multiply(result, {scale, 0});
}
}
if (result.first > result.second) {
std::swap(result.first, result.second);
}
return result;
}
/// @brief Count integer pairs (a,b), including signs and order, satisfying
/// a^2 + b^2 = n.
inline unsigned __int128 sum_two_squares_representation_count(std::uint64_t n) {
if (n == 0) {
return 1;
}
unsigned __int128 result = 4;
for (auto [prime, exponent] : factorize(n)) {
if (prime % 4 == 3 && exponent % 2 == 1) {
return 0;
}
if (prime % 4 == 1) {
result *= exponent + 1;
}
}
return result;
}
/// @brief Return every ordered non-negative pair (a,b) with a^2 + b^2 = n.
/// The result is sorted and contains no duplicates. Requires n <= 1e18.
inline std::vector<std::pair<std::uint64_t, std::uint64_t>>
all_sum_two_squares(std::uint64_t n) {
using namespace sum_two_squares_internal;
assert(n <= 1000000000000000000ULL);
if (n == 0) {
return {{0, 0}};
}
auto factors = factorize(n);
for (auto [prime, exponent] : factors) {
if (prime % 4 == 3 && exponent % 2 == 1) {
return {};
}
}
std::vector<gaussian> representatives{{1, 0}};
for (auto [prime, exponent] : factors) {
if (prime % 4 == 3) {
i128 scale = 1;
for (int i = 0; i < exponent / 2; i++) {
scale *= prime;
}
for (auto &[real, imaginary] : representatives) {
real *= scale;
imaginary *= scale;
}
continue;
}
gaussian prime_factor;
if (prime == 2) {
prime_factor = {1, 1};
gaussian multiplier = signed_power(prime_factor, exponent);
for (auto &value : representatives) {
value = signed_multiply(value, multiplier);
}
continue;
}
auto representation = prime_representation(prime);
assert(representation.has_value());
prime_factor = {representation->first, representation->second};
std::vector<gaussian> powers(exponent + 1, {1, 0});
for (int i = 0; i < exponent; i++) {
powers[i + 1] = signed_multiply(powers[i], prime_factor);
}
std::vector<gaussian> next;
next.reserve(representatives.size() * (exponent + 1));
for (gaussian current : representatives) {
for (int chosen = 0; chosen <= exponent; chosen++) {
gaussian conjugate = powers[exponent - chosen];
conjugate.second = -conjugate.second;
gaussian factor = signed_multiply(powers[chosen], conjugate);
next.push_back(signed_multiply(current, factor));
}
}
representatives.swap(next);
}
std::vector<std::pair<u64, u64>> result;
for (auto [real, imaginary] : representatives) {
while (real <= 0 || imaginary < 0) {
i128 old_real = real;
real = -imaginary;
imaginary = old_real;
}
result.emplace_back(u64(real), u64(imaginary));
if (imaginary == 0) {
result.emplace_back(0, u64(real));
}
}
std::sort(result.begin(), result.end());
result.erase(std::unique(result.begin(), result.end()), result.end());
return result;
}
} // namespace noya