mod 平方根

$a \not\equiv 0$ のとき、 $x^2 \equiv a \pmod{p}$ に解があるのは $a^{(p-1)/2} \equiv 1 \pmod{p}$ のときに限ります（オイラーの規準）。 $p - 1 = 2^s q$（ $q$ は奇数）と表し、平方非剰余 $z$ を 1 つ見つけておきます。 $x = a^{(q+1)/2}$, $t = a^q$ から始めると $x^2 \equiv a t$ で、 $t$ の位数は $2$ の冪です。 $z^q$ の冪を掛けて $x^2 \equiv a t$ を保ったまま $t$ の位数を小さくしていき、 $t = 1$ になったときの $x$ が答えです（Tonelli–Shanks のアルゴリズム）。

• mod_sqrt(a, p): $x^2 \equiv a \pmod{p}$ を満たす $x$（ $0 \le x < p$）を 1 つ返す。なければ -1。 $p$ は素数（ $p = 2$ でもよい）。 $O(\log^2 p)$（ $z$ を探す部分は期待値）

もう 1 つの解は $p - x$ です（ $a \equiv 0$ のときは $x = 0$ だけ）。 $p \equiv 3 \pmod{4}$ なら $s = 1$ で、 $x = a^{(p+1)/4}$ がそのまま答えになります。

コード

def mod_sqrt(a, p):
    # x^2 ≡ a (mod p) となる x を 1 つ（p は素数、なければ -1）。もう 1 つの解は p - x
    a %= p
    if a < 2 or p == 2:
        return a
    if pow(a, (p - 1) // 2, p) != 1:
        return -1  # 平方非剰余（オイラーの規準）
    q, s = p - 1, 0
    while q % 2 == 0:
        q //= 2
        s += 1
    z = 2
    while pow(z, (p - 1) // 2, p) == 1:
        z += 1  # 平方非剰余 z を探す
    c = pow(z, q, p)  # 位数がちょうど 2^s
    x = pow(a, (q + 1) // 2, p)
    t = pow(a, q, p)  # x^2 ≡ a t を保ち、t を 1 にする（Tonelli–Shanks）
    while t != 1:
        i, u = 0, t
        while u != 1:
            u = u * u % p
            i += 1  # t の位数は 2^i
        b = pow(c, 1 << (s - i - 1), p)
        s = i
        c = b * b % p
        x = x * b % p
        t = t * c % p
    return x