Math/tests/sqrt_mod.test.cpp

• View this file on GitHub
• Last update: 2022-12-25 00:11:23+08:00
• Problem: https://judge.yosupo.jp/problem/sqrt_mod

Depends on

• Math/NumberTheory/SqrtMod.h

Code

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

#include <bits/stdc++.h>
using namespace std;

// for 64-bit, use mt19937_64
mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
long long get_rand(long long r) {
    return uniform_int_distribution<long long> (0, r-1)(rng);
}

#include "../NumberTheory/SqrtMod.h"

int32_t main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int ntest; cin >> ntest;
    while (ntest--) {
        long long n, p; cin >> n >> p;
        long long res = sqrtMod(n, p);
        cout << res << '\n';
    }
    return 0;
}

#line 1 "Math/tests/sqrt_mod.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/sqrt_mod"

#include <bits/stdc++.h>
using namespace std;

// for 64-bit, use mt19937_64
mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
long long get_rand(long long r) {
    return uniform_int_distribution<long long> (0, r-1)(rng);
}

#line 1 "Math/NumberTheory/SqrtMod.h"
// sqrtMod(X, q), with q is prime, returns:
// a where a*a = X
// -1 if no solution
//
// Note:
// - there are either 1 or 2 solutions, a and p - a (which can be same).
//   This code return smaller solution
//
// Copied from https://judge.yosupo.jp/submission/59210
//
// Tested:
// - (p <= 10^9) https://judge.yosupo.jp/problem/sqrt_mod
// - (print all sols) https://oj.vnoi.info/problem/jacobi
// - https://oj.vnoi.info/problem/newj

// sqrtMod {{{
using ll = long long;
ll euclid(ll x, ll y, ll &k, ll &l) {
    if (y == 0) {
        k = 1;
        l = 0;
        return x;
    }
    ll g = euclid(y, x % y, l, k);
    l -= k * (x / y);
    return g;
}

ll mult(ll x, ll y, ll md) {
    return x * y % md;
}
ll bin_pow(ll x, ll p, ll md) {
    if (p == 0) return 1;
    if (p & 1) return mult(x, bin_pow(x, p - 1, md), md);
    return bin_pow(mult(x, x, md), p / 2, md);
}

ll Cipolla(ll X, ll q) {
    ll pw = (q - 1) / 2;
    int K = 60;
    while((1LL << K) > pw) K--;
    while(true) {
        ll t = get_rand(q);
        ll a = 0, b = 0, c = 1;
        for (int k = K; k >= 0; k--) {
            a = b * b % q;
            b = 2 * b * c % q;
            c = (c * c + a * X) % q;
            if (((pw >> k) & 1) == 0) continue;
            a = b;
            b = (b * t + c) % q;
            c = (c * t + a * X) % q;
        }
        if (b == 0) continue;
        c = (c + q - 1) % q;
        ll k, l;
        euclid(b, q, k, l);
        c = -c * k % q;
        if (c < 0) c += q;
        if (c * c % q == X) return c;
    }
    assert(false);
}

ll sqrtMod(ll X, ll q) {
    X %= q;
    if (q == 2 || X == 0) return min(X, q-X);
    if (bin_pow(X, (q - 1) / 2, q) != 1) return -1;
    if (q % 4 == 3) {
        ll res = bin_pow(X, (q + 1) / 4, q);
        return min(res, q - res);
    }
    auto res = (Cipolla(X, q) % q + q) % q;
    return min(res, q-res);
}
// }}}
#line 13 "Math/tests/sqrt_mod.test.cpp"

int32_t main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int ntest; cin >> ntest;
    while (ntest--) {
        long long n, p; cin >> n >> p;
        long long res = sqrtMod(n, p);
        cout << res << '\n';
    }
    return 0;
}

Back to top page