Z/pZ上の平方根(Lehmerのアルゴリズム)

目的

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

計算量

一回の試行の計算量が$O(\log p)$。
試行は確率$1-(p-1)/2p$で失敗する。
成功するまで試行を繰り返す。

使い方

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(2, 1);
    p2.add(1, -2 * b);
    p2.add(0, b * b - a);
    p1.add(1, 1);
    p1 = p1.pow((mod - 1) / 2, p2);
    p1.add(0, -1);
    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) {
      ret.add(i, val[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) {
      modPoly.add(i, p.val[i] * rev);
    }
    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);
    ret.add(0, 1);
    Poly pow = new Poly(mod);
    for (int i = 0; i < val.length; ++i) {
      pow.add(i, val[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 夏休みの思い出(2)