Skip to main content
2 of 8
added 17 characters in body; edited tags
Jamal
  • 34.9k
  • 13
  • 133
  • 237

Python implementation of prime modular square root

I did an implementation of the Tonelli-Shanks algorithm as defined on Wikipedia. I put it here for review and sharing purpose.

Legendre Symbol implementation:

def legendre_symbol(a, p):
    """
    Legendre symbol
    Define if a is a quadratic residue modulo odd prime
    http://en.wikipedia.org/wiki/Legendre_symbol
    """
    ls = pow(a, (p - 1)/2, p)
    if ls == p - 1:
        return -1
    return ls

Prime modular square root (I just renamed the solution variable R to x):

def prime_mod_sqrt(a, p):
    """
    Square root modulo prime number
    Solve the equation
        x^2 = a mod p
    and return list of x solution
    http://en.wikipedia.org/wiki/Tonelli-Shanks_algorithm
    """
    a %= p

    # Simple case
    if a == 0:
        return [0]
    if p == 2:
        return [a]

    # Check solution existence on odd prime
    if legendre_symbol(a, p) != 1:
        return []

    # Simple case
    if p % 4 == 3:
        x = pow(a, (p + 1)/4, p)
        return [x, p-x]

    # Factor p-1 on the form q * 2^s (with Q odd)
    q, s = p - 1, 0
    while q % 2 == 0:
        s += 1
        q //= 2

    # Select a z which is a quadratic non resudue modulo p
    z = 1
    while legendre_symbol(z, p) != -1:
        z += 1
    c = pow(z, q, p)

    # Search for a solution
    x = pow(a, (q + 1)/2, p)
    t = pow(a, q, p)
    m = s
    while t != 1:
        for i in xrange(1, m):
            if pow(t, 2**i, p) == 1:
                b = pow(c, 2**(m-i-1), p)
                x = (x * b) % p
                t = (t * b * b) % p
                m = i
                break

    return [x, p-x]

If you have any optimization or found any error, please report it.

Phong
  • 263
  • 3
  • 10