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

目的

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

使い方

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

ソースコード

static long sqrt(long a, long p) {
    if (a == 0)
	return 0;
    if (a == 1)
	return 1;
    long s = 0;
    long tmp = p - 1;
    while (tmp % 2 == 0) {
	++s;
	tmpfont color=#006636> /= 2;
    }
    long Q = tmp;
    long z = 1;
    while (pow(z, (p - 1) / 2, p) != p - 1) {
	++z;
    }
    long M = s;
    long c = pow(z, Q, p);
    long t = pow(a, Q, p);
    long R = pow(a, (Q + 1) / 2, p);
    while (t != 1) {
	//if (pow(t, pow(2, M - 1, p - 1), p) != 1)
	//throw new AssertionError();
	//if (pow(c, pow(2, M - 1, p - 1), p) != p - 1)
	//throw new AssertionError(M);
	long cur = t;
	int i;
	for (i = 1; i < M; ++i) {
	    cur = cur * cur % p;
	    if (cur == 1)
		break;
	}
	if (cur != 1) {
	    throw new AssertionError();
	}
	long b = pow(c, pow(2, M - i - 1, p - 1), p);
	M = i;
	c = b * b % p;
	t = t * b % p * b % p;
	R = R * b % p;
    }
    if (R * R % p != a) {
	throw new AssertionError();
    }
    return R;
}


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)