Merge pull request #414 from hitonanode/gaussian-integer

Gaussian integer

Author: hitonanode

.github/workflows/verify.yml

@@ -14,13 +14,15 @@ jobs:
     - uses: actions/checkout@v1
 
     - name: Set up Python
-      uses: actions/setup-python@v1
-    
+      uses: actions/setup-python@v5
+
     - name: Install boost
       run: sudo apt install -y libboost-dev
 
     - name: Install dependencies
-      run: pip3 install -U online-judge-verify-helper
+      run: |
+        python -m pip install -U pip
+        python -m pip install -U "setuptools<81" online-judge-verify-helper
 
     - name: Clone AC Library
       run: cd ${GITHUB_WORKSPACE}/../ && git clone https://github.com/atcoder/ac-library.git

number/gaussian_integer.hpp

@@ -0,0 +1,81 @@
+#pragma once
+#include <cassert>
+#include <utility>
+
+struct GaussianInteger {
+    using Int = long long;
+    using G = GaussianInteger;
+    Int x, y; // x + yi
+
+    Int re() const noexcept { return x; }
+    Int im() const noexcept { return y; }
+
+    explicit GaussianInteger(Int x_ = 0, Int y_ = 0) : x(x_), y(y_) {}
+
+    G &operator+=(const G &n) {
+        x += n.x, y += n.y;
+        return *this;
+    }
+    G &operator-=(const G &n) {
+        x -= n.x, y -= n.y;
+        return *this;
+    }
+    G &operator*=(const G &n) {
+        const Int nx = x * n.x - y * n.y, ny = x * n.y + y * n.x;
+        x = nx, y = ny;
+        return *this;
+    }
+
+    // Euclidean division
+    G &operator/=(const G &n) {
+        const Int d = n.norm();
+        assert(d != 0);
+        const Int nx = x * n.x + y * n.y, ny = y * n.x - x * n.y;
+        auto div_round = [](Int num, Int den) {
+            return (num >= 0) ? (num + den / 2) / den : (num - den / 2) / den;
+        };
+        x = div_round(nx, d), y = div_round(ny, d);
+        return *this;
+    }
+
+    G &operator%=(const G &n) {
+        *this -= (*this / n) * n;
+        return *this;
+    }
+
+    G operator+(const G &n) const { return G(*this) += n; }
+    G operator-(const G &n) const { return G(*this) -= n; }
+    G operator*(const G &n) const { return G(*this) *= n; }
+    G operator/(const G &n) const { return G(*this) /= n; }
+    G operator%(const G &n) const { return G(*this) %= n; }
+
+    Int norm() const { return x * x + y * y; }
+
+    G conj() const { return G{x, -y}; }
+
+    static G gcd(G a, G b) {
+        while (b.x != 0 or b.y != 0) {
+            a %= b;
+            std::swap(a, b);
+        }
+        return a.canonical();
+    }
+    friend G gcd(G a, G b) { return G::gcd(a, b); }
+
+    bool operator==(const G &n) const { return x == n.x and y == n.y; }
+    bool operator!=(const G &n) const { return !(*this == n); }
+
+    template <class OStream> friend OStream &operator<<(OStream &os, const G &g) {
+        return os << g.x << (g.y >= 0 ? "+" : "") << g.y << "i";
+    }
+
+    // move to first quadrant (x > 0, y >= 0)
+    G canonical() const {
+        if (x > 0 and y >= 0) return *this;
+        if (x <= 0 and y > 0) return G{y, -x};
+        if (x < 0 and y <= 0) return G{-x, -y};
+        return G{-y, x};
+    }
+
+    std::pair<Int, Int> to_pair() const { return {x, y}; }
+};

number/gaussian_integer.md

@@ -0,0 +1,10 @@
+---
+title: Gaussian Integer （ガウス整数）
+documentation_of: ./gaussian_integer.hpp
+---
+
+ガウス整数を扱う．
+
+## 問題例
+
+- [Library Checker: Gcd of Gaussian Integers](https://judge.yosupo.jp/problem/gcd_of_gaussian_integers)

number/sqrt_mod.hpp

@@ -0,0 +1,47 @@
+#pragma once
+#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);
+}

number/sqrt_mod.md

@@ -0,0 +1,26 @@
+---
+title: Square root modulo prime （平方剰余）
+documentation_of: ./sqrt_mod.hpp
+---
+
+素数 $p$ と整数 $a$ に対して，$x^2 \equiv a \pmod{p}$ を満たす $x$ を求める．解が存在しない場合は $-1$ を返す．Tonelli-Shanks algorithm に基づく．
+
+## 使用方法
+
+```cpp
+int a, p;
+int x = sqrt_mod<int>(a, p);  // x^2 ≡ a (mod p), or -1
+
+long long al, pl;
+long long xl = sqrt_mod<long long>(al, pl);  // __int128 を内部で使用
+```
+
+`Int` が `int` のとき内部で `long long`，`long long` のとき `__int128` を乗算のオーバーフロー回避に使用する．
+
+## 問題例
+
+- [Library Checker: Sqrt Mod](https://judge.yosupo.jp/problem/sqrt_mod)
+
+## リンク
+
+- [Tonelli–Shanks algorithm - Wikipedia](https://en.wikipedia.org/wiki/Tonelli%E2%80%93Shanks_algorithm)

number/square_sums.hpp

@@ -0,0 +1,19 @@
+#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));
+}

number/square_sums.md

@@ -0,0 +1,21 @@
+---
+title: Sum of two squares （二つの平方数の和）
+documentation_of: ./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)$ を返す．
+
+## 使用方法
+
+```cpp
+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](https://en.wikipedia.org/wiki/Fermat%27s_theorem_on_sums_of_two_squares)

number/test/gcd_of_gaussian_integers.test.cpp

@@ -0,0 +1,20 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/gcd_of_gaussian_integers"
+
+#include "../gaussian_integer.hpp"
+
+#include <iostream>
+using namespace std;
+
+int main() {
+    cin.tie(nullptr), ios::sync_with_stdio(false);
+
+    int T;
+    cin >> T;
+    while (T--) {
+        long long a, b, c, d;
+        cin >> a >> b >> c >> d;
+        GaussianInteger g1(a, b), g2(c, d);
+        const auto g = gcd(g1, g2);
+        cout << g.x << " " << g.y << "\n";
+    }
+}

number/test/sqrt_mod.test.cpp

@@ -0,0 +1,15 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/sqrt_mod"
+#include "../sqrt_mod.hpp"
+
+#include <iostream>
+using namespace std;
+
+int main() {
+    int T;
+    cin >> T;
+    while (T--) {
+        int Y, P;
+        cin >> Y >> P;
+        cout << sqrt_mod<int>(Y, P) << '\n';
+    }
+}