Sum of two squares （二つの平方数の和）
(number/square_sums.hpp)

• View this file on GitHub
• Last update: 2026-03-01 20:50:35+09:00
• Include: #include "number/square_sums.hpp"

素数 $p$ に対して，$x^2 + y^2 = p$ を満たす非負整数の組 $(x, y)$ （$x \le y$）を求める．Fermat の二平方和定理に基づく．

$p = 2$ のとき $(1, 1)$ を，$p \equiv 1 \pmod{4}$ のとき解を返す．$p \equiv 3 \pmod{4}$ のときは解が存在せず $(-1, -1)$ を返す．

使用方法

int p;  // prime
auto [x, y] = SumOfTwoSquaresPrime<int>(p);  // x <= y, x^2 + y^2 = p

内部で sqrt_mod を呼び出し，$-1$ の平方根 $r$ を求めた後，$p$ と $r$ に対するユークリッドの互除法を $r_i \le \sqrt{p}$ となるまで適用する．

リンク

• Fermat’s theorem on sums of two squares - Wikipedia

Depends on

• Square root modulo prime （平方剰余） (number/sqrt_mod.hpp)

Code

#pragma once
#include <cmath>

#include "sqrt_mod.hpp"

// Fermat's theorem on sums of two squares
// Solve x^2 + y^2 = p for prime p, where p != 3 (mod 4)
template <class Int> std::pair<Int, Int> SumOfTwoSquaresPrime(Int p) {
    if (p == 2) return {1, 1};
    if (p % 4 != 1) return {-1, -1};

    Int r0 = p, r1 = sqrt_mod<Int>(p - 1, p);
    for (const Int limit = std::sqrtl(p); r1 > limit;) {
        Int next_r = r0 % r1;
        r0 = r1, r1 = next_r;
    }

    return std::minmax(r1, (Int)std::sqrtl(p - r1 * r1));
}

#line 2 "number/square_sums.hpp"
#include <cmath>

#line 2 "number/sqrt_mod.hpp"
#include <algorithm>
#include <type_traits>

// Solve x^2 == a (MOD p) for prime p
// If no solution exists, return -1
template <class Int> Int sqrt_mod(Int a, Int p) {
    using Long =
        std::conditional_t<std::is_same_v<Int, int>, long long,
                           std::conditional_t<std::is_same_v<Int, long long>, __int128, void>>;

    auto pow = [&](Int x, long long n) {
        Int ans = 1, tmp = x;
        while (n) {
            if (n & 1) ans = (Long)ans * tmp % p;
            tmp = (Long)tmp * tmp % p, n /= 2;
        }
        return ans;
    };
    if (a == 0) return 0;

    a = (a % p + p) % p;
    if (p == 2) return a;
    if (pow(a, (p - 1) / 2) != 1) return -1;

    int b = 1;
    while (pow(b, (p - 1) / 2) == 1) ++b;

    int e = 0;
    Int m = p - 1;
    while (m % 2 == 0) m /= 2, ++e;

    Int x = pow(a, (m - 1) / 2), y = (Long)x * x % p * a % p;
    x = (Long)x * a % p;
    Int z = pow(b, m);
    while (y != 1) {
        int j = 0;
        Int t = y;
        while (t != 1) t = (Long)t * t % p, ++j;
        z = pow(z, 1LL << (e - j - 1));
        x = (Long)x * z % p;
        z = (Long)z * z % p;
        y = (Long)y * z % p;
        e = j;
    }
    return std::min(x, p - x);
}
#line 5 "number/square_sums.hpp"

// Fermat's theorem on sums of two squares
// Solve x^2 + y^2 = p for prime p, where p != 3 (mod 4)
template <class Int> std::pair<Int, Int> SumOfTwoSquaresPrime(Int p) {
    if (p == 2) return {1, 1};
    if (p % 4 != 1) return {-1, -1};

    Int r0 = p, r1 = sqrt_mod<Int>(p - 1, p);
    for (const Int limit = std::sqrtl(p); r1 > limit;) {
        Int next_r = r0 % r1;
        r0 = r1, r1 = next_r;
    }

    return std::minmax(r1, (Int)std::sqrtl(p - r1 * r1));
}

Back to top page