数学/平方剰余
(test/math/mod_sqrt.test.cpp)

• View this file on GitHub
• Last update: 2023-02-25 16:35:06+09:00
• Problem: https://judge.yosupo.jp/problem/sqrt_mod

Depends on

• 繰り返し二乗法 / 二分累乗法 / バイナリ法 (include/emthrm/math/mod_pow.hpp)
• 平方剰余 (quadratic residue) (include/emthrm/math/mod_sqrt.hpp)

Code

/*
 * @title 数学/平方剰余
 *
 * verification-helper: PROBLEM https://judge.yosupo.jp/problem/sqrt_mod
 */

#include <iostream>

#include "emthrm/math/mod_sqrt.hpp"

int main() {
  int t;
  std::cin >> t;
  while (t--) {
    int y, p;
    std::cin >> y >> p;
    std::cout << emthrm::mod_sqrt(y, p) << '\n';
  }
  return 0;
}

#line 1 "test/math/mod_sqrt.test.cpp"
/*
 * @title 数学/平方剰余
 *
 * verification-helper: PROBLEM https://judge.yosupo.jp/problem/sqrt_mod
 */

#include <iostream>

#line 1 "include/emthrm/math/mod_sqrt.hpp"



#include <random>

#line 1 "include/emthrm/math/mod_pow.hpp"



namespace emthrm {

long long mod_pow(long long x, long long n, const int m) {
  if ((x %= m) < 0) x += m;
  long long res = 1;
  for (; n > 0; n >>= 1) {
    if (n & 1) res = (res * x) % m;
    x = (x * x) % m;
  }
  return res;
}

}  // namespace emthrm


#line 7 "include/emthrm/math/mod_sqrt.hpp"

namespace emthrm {

long long mod_sqrt(long long a, const int p) {
  if ((a %= p) < 0) a += p;
  if (a == 0) [[unlikely]] return 0;
  if (p == 2) [[unlikely]] return 1;
  if (mod_pow(a, (p - 1) >> 1, p) == p - 1) return -1;
  if (p % 4 == 3) return mod_pow(a, (p + 1) >> 2, p);
  int s = 1, q = (p - 1) >> 1;
  for (; !(q & 1); q >>= 1) {
    ++s;
  }
  static std::mt19937_64 engine(std::random_device {} ());
  std::uniform_int_distribution<> dist(2, p - 1);
  long long z;
  do {
    z = dist(engine);
  } while (mod_pow(z, (p - 1) >> 1, p) == 1);
  int m = s;
  long long c = mod_pow(z, q, p), r = mod_pow(a, (q - 1) >> 1, p);
  long long t = a * r % p * r % p;
  r = (r * a) % p;
  while (t != 1) {
    long long t2 = t * t % p;
    for (int i = 1; i < m; ++i) {
      if (t2 == 1) {
        const long long b = mod_pow(c, 1 << (m - i - 1), p);
        m = i;
        r = (r * b) % p;
        c = b * b % p;
        t = (t * c) % p;
        break;
      }
      t2 = (t2 * t2) % p;
    }
  }
  return r;
}

}  // namespace emthrm


#line 10 "test/math/mod_sqrt.test.cpp"

int main() {
  int t;
  std::cin >> t;
  while (t--) {
    int y, p;
    std::cin >> y >> p;
    std::cout << emthrm::mod_sqrt(y, p) << '\n';
  }
  return 0;
}

Back to top page