Mod sqrt
(Mylib/Number/Mod/mod_sqrt.cpp)

• View this file on GitHub
• Last update: 2021-04-23 23:44:44+09:00
• Link: View error logs on GitHub Actions

Operations

• mod_sqrt(a, p)
  • $x ^ 2 = a \pmod p$を満たすxを求める。

Requirements

Notes

Problems

References

• https://ja.wikipedia.org/wiki/%E5%B9%B3%E6%96%B9%E5%89%B0%E4%BD%99%E3%81%AE%E7%9B%B8%E4%BA%92%E6%B3%95%E5%89%87
• https://qiita.com/drken/items/3b4fdf0a78e7a138cd9a
• https://en.wikipedia.org/wiki/Tonelli%E2%80%93Shanks_algorithm

Depends on

• Mod pow (Mylib/Number/Mod/mod_pow.cpp)

Required by

• Formal power series (Sqrt) (Mylib/Math/fps_sqrt.cpp)

Verified with

• test/yosupo-judge/sqrt_mod/main.test.cpp
• test/yosupo-judge/sqrt_of_formal_power_series/main.test.cpp

Code

#pragma once
#include <optional>
#include <random>
#include "Mylib/Number/Mod/mod_pow.cpp"

namespace haar_lib {
  std::optional<int64_t> mod_sqrt(int64_t a, int64_t p) {
    if (p == 2) return a % 2;
    if (a == 0) return 0;

    int64_t b = mod_pow(a, (p - 1) / 2, p);

    if (b == p - 1) return {};
    if (p % 4 == 3) return mod_pow(a, (p + 1) / 4, p);

    int64_t q = p - 1, s = 0;
    while (q % 2 == 0) {
      q /= 2;
      s += 1;
    }

    static std::mt19937_64 rand(time(0));
    std::uniform_int_distribution<> dist(0, p - 1);

    int64_t z;
    while (1) {
      z = dist(rand);
      if (mod_pow(z, (p - 1) / 2, p) == p - 1) break;
    }

    int64_t m = s;
    int64_t c = mod_pow(z, q, p);
    int64_t t = mod_pow(a, q, p);
    int64_t r = mod_pow(a, (q + 1) / 2, p);

    while (1) {
      if (t == 0) return 0;
      if (t == 1) return r;

      int i = 1;
      for (int64_t T = t; i < m; ++i) {
        (T *= T) %= p;
        if (T == 1) break;
      }

      int64_t b = mod_pow(c, 1LL << (m - i - 1), p);

      m = i;
      c = b * b % p;
      (t *= b * b % p) %= p;
      (r *= b) %= p;
    }
  }
}  // namespace haar_lib

#line 2 "Mylib/Number/Mod/mod_sqrt.cpp"
#include <optional>
#include <random>
#line 2 "Mylib/Number/Mod/mod_pow.cpp"
#include <cstdint>

namespace haar_lib {
  constexpr int64_t mod_pow(int64_t n, int64_t p, int64_t m) {
    int64_t ret = 1;
    while (p > 0) {
      if (p & 1) (ret *= n) %= m;
      (n *= n) %= m;
      p >>= 1;
    }
    return ret;
  }
}  // namespace haar_lib
#line 5 "Mylib/Number/Mod/mod_sqrt.cpp"

namespace haar_lib {
  std::optional<int64_t> mod_sqrt(int64_t a, int64_t p) {
    if (p == 2) return a % 2;
    if (a == 0) return 0;

    int64_t b = mod_pow(a, (p - 1) / 2, p);

    if (b == p - 1) return {};
    if (p % 4 == 3) return mod_pow(a, (p + 1) / 4, p);

    int64_t q = p - 1, s = 0;
    while (q % 2 == 0) {
      q /= 2;
      s += 1;
    }

    static std::mt19937_64 rand(time(0));
    std::uniform_int_distribution<> dist(0, p - 1);

    int64_t z;
    while (1) {
      z = dist(rand);
      if (mod_pow(z, (p - 1) / 2, p) == p - 1) break;
    }

    int64_t m = s;
    int64_t c = mod_pow(z, q, p);
    int64_t t = mod_pow(a, q, p);
    int64_t r = mod_pow(a, (q + 1) / 2, p);

    while (1) {
      if (t == 0) return 0;
      if (t == 1) return r;

      int i = 1;
      for (int64_t T = t; i < m; ++i) {
        (T *= T) %= p;
        if (T == 1) break;
      }

      int64_t b = mod_pow(c, 1LL << (m - i - 1), p);

      m = i;
      c = b * b % p;
      (t *= b * b % p) %= p;
      (r *= b) %= p;
    }
  }
}  // namespace haar_lib

Back to top page