﻿ Z/pZ上の平方根（Lehmerのアルゴリズム）

# Z/pZ上の平方根（Lehmerのアルゴリズム）

## 目的

$x^2 = a \mod p$を満たす整数$x$をLehmerのアルゴリズムによって求めます。

## 使い方

long sqrt(long a,long p)
$x^2 = a \mod p$を満たす$x$をreturnします。

## 参考文献

Richard M. Karp. An introduction to randomized algorithms. Discrete Applied Mathematics, Vol. 34, No. 1-3, pp. 165–201, 1991.

## ソースコード

static long sqrt(long a, long mod) {
for (int b = 0; b < 100; ++b) {
Poly p1 = new Poly(mod);
Poly p2 = new Poly(mod);// (x-b)^2-a=0 mod p
p2.add(0, b * b - a);
p1 = p1.pow((mod - 1) / 2, p2);
Poly gcd = p1.gcd(p2);
if (gcd.deg() == 1) {
long ans = (-b - inv(gcd.val[1], mod) * gcd.val[0]) % mod;
while (ans < 0)
ans += mod;
return ans;
}
}
throw new AssertionError();
}

static class Poly {
long[] val = new long[3];
long mod;

public Poly(long mod) {
this.mod = mod;
}

void add(int deg, long v) {
val[deg] += v;
while (val[deg] < 0)
val[deg] += mod;
val[deg] %= mod;
}

int deg() {
int deg = val.length - 1;
while (deg > 0 && val[deg] == 0)
--deg;
return deg;
}

Poly mul(Poly p) {
Poly ret = new Poly(mod);
for (int i = 0; i < this.val.length; ++i) {
if (val[i] == 0)
continue;
for (int j = 0; j < p.val.length; ++j) {
if (p.val[j] == 0)
continue;
ret.add(i + j, this.val[i] * p.val[j]);
}
}
return ret;
}

Poly mod(Poly p) {
Poly ret = new Poly(mod);
Poly modPoly = new Poly(mod);
for (int i = 0; i < val.length; ++i) {
}
int max = p.deg();
if (max == 0 && p.val[max] == 0)
return new Poly(mod);
long rev = (long) inv(p.val[max], mod);
for (int i = 0; i < p.val.length; ++i) {
}
for (int i = ret.val.length - 1; i >= max; --i) {
if (ret.val[i] == 0)
continue;
long tmp = ret.val[i];
for (int j = 0; j <= max; ++j) {
ret.add(i - j, -tmp % mod * modPoly.val[max - j] % mod);
}
}
return ret;
}

Poly gcd(Poly p) {
int deg1 = this.deg();
int deg2 = p.deg();
if (deg1 > deg2) {
return p.gcd(this);
}
if (deg1 == 0 && val[deg1] == 0)
return p;
return p.mod(this).gcd(this);
}

Poly pow(long n, Poly polyMod) {
int deg = this.deg();
if (val[deg] == 0)
return new Poly(mod);
Poly ret = new Poly(mod);
Poly pow = new Poly(mod);
for (int i = 0; i < val.length; ++i) {
}
for (; n > 0; n >>= 1, pow = pow.mul(pow), pow = pow.mod(polyMod)) {
if (n % 2 == 1) {
ret = ret.mul(pow);
ret = ret.mod(polyMod);
}
}
return ret;
}
}

public static long inv(long a, long mod) {
long b = mod;
long p = 1, q = 0;
while (b > 0) {
long c = a / b;
long d;
d = a;
a = b;
b = d % b;
d = p;
p = q;
q = d - c * q;
}
return p < 0 ? p + mod : p;
}

static long pow(long a, long n, long mod) {
long ret = 1;
for (; n > 0; n >>= 1, a = (a * a) % mod) {
if (n % 2 == 1) {
ret = (ret * a) % mod;
}
}
return ret;
}

## Verified

yukicoder No.551 夏休みの思い出（２）