Computing modular square roots in Python

March 7th, 2009 at 11:59 am

Consider the congruence of the form:

$$!x^2 \equiv n \pmod p$$

n is a quadradic residue (mod n). What is x? In normal arithmetic, this is equivalent to finding the square root of a number. In modular arithmetic, x is the modular square root of n modulo p.

Now, in the general case, this is a very difficult problem to solve. In fact, it’s equivalent to integer factorization, because no efficient algorithm is known to find the modular square root modulo a composite number, and if the modulo is composite it has to be factored first.

But when p is prime, an efficient polynomial algorithm exists for computing x. This is the Tonelli-Shanks algorithm.

Computing modular square roots is probably not one of those things you do daily, but I ran into it while solving a Project Euler problem. So I’m posting the Python implementation of the Tonelli-Shanks algorithm here. It is based on the explanation in the paper "Square roots from 1; 24, 51, 10 to Dan Shanks" by Ezra Brown, as I found the Wikipedia algorithm hard to follow.

The code is tested, and as far as I can tell works correctly and efficiently:

def modular_sqrt(a, p):
    """ Find a quadratic residue (mod p) of 'a'. p
        must be an odd prime.

        Solve the congruence of the form:
            x^2 = a (mod p)
        And returns x. Note that p - x is also a root.

        0 is returned is no square root exists for
        these a and p.

        The Tonelli-Shanks algorithm is used (except
        for some simple cases in which the solution
        is known from an identity). This algorithm
        runs in polynomial time (unless the
        generalized Riemann hypothesis is false).
    """
    # Simple cases
    #
    if legendre_symbol(a, p) != 1:
        return 0
    elif a == 0:
        return 0
    elif p == 2:
        return p
    elif p % 4 == 3:
        return pow(a, (p + 1) / 4, p)

    # Partition p-1 to s * 2^e for an odd s (i.e.
    # reduce all the powers of 2 from p-1)
    #
    s = p - 1
    e = 0
    while s % 2 == 0:
        s /= 2
        e += 1

    # Find some 'n' with a legendre symbol n|p = -1.
    # Shouldn't take long.
    #
    n = 2
    while legendre_symbol(n, p) != -1:
        n += 1

    # Here be dragons!
    # Read the paper "Square roots from 1; 24, 51,
    # 10 to Dan Shanks" by Ezra Brown for more
    # information
    #

    # x is a guess of the square root that gets better
    # with each iteration.
    # b is the "fudge factor" - by how much we're off
    # with the guess. The invariant x^2 = ab (mod p)
    # is maintained throughout the loop.
    # g is used for successive powers of n to update
    # both a and b
    # r is the exponent - decreases with each update
    #
    x = pow(a, (s + 1) / 2, p)
    b = pow(a, s, p)
    g = pow(n, s, p)
    r = e

    while True:
        t = b
        m = 0
        for m in xrange(r):
            if t == 1:
                break
            t = pow(t, 2, p)

        if m == 0:
            return x

        gs = pow(g, 2 ** (r - m - 1), p)
        g = (gs * gs) % p
        x = (x * gs) % p
        b = (b * g) % p
        r = m


def legendre_symbol(a, p):
    """ Compute the Legendre symbol a|p using
        Euler's criterion. p is a prime, a is
        relatively prime to p (if p divides
        a, then a|p = 0)

        Returns 1 if a has a square root modulo
        p, -1 otherwise.
    """
    ls = pow(a, (p - 1) / 2, p)
    return -1 if ls == p - 1 else ls

Related posts:

  1. Efficient modular exponentiation algorithms
  2. more thoughts on previous problem
  3. back to the roots
  4. Distributed computing in Python with multiprocessing
  5. Project Euler problem 66 and continued fractions

13 Responses to “Computing modular square roots in Python”

  1. SteveNo Gravatar Says:

    The lines
    elif p == 2:
    return n

    look wrong to me – ‘n’ has not been defined.

    Also, are all the divisions integer divisions (that is, do they all produce integer results, and could you use the ‘//’ operator) ?

  2. Nik SkaNo Gravatar Says:

    Thank you!

  3. bobNo Gravatar Says:

    You probably want to call it “Tonelli-Shanks.” You refer to it as Snanks-Tonelli, and Shaks-Tonelli.

  4. elibenNo Gravatar Says:

    bob,

    Fixed, thanks.

  5. SethNo Gravatar Says:

    I’ve read this algorithm is easy to modify for cube roots, I don’t suppose you know how…?

  6. elibenNo Gravatar Says:

    Seth, I don’t :-)

  7. GregSNo Gravatar Says:

    Steve’s bug still remains, minor as it is. replace

    elif p == 2:
        return n

    with

    elif p == 2:
       return p
  8. elibenNo Gravatar Says:

    Steve & GregS,

    Yep, this is a typo. Fixed to return p. Thanks!

    Steve – yes, these are integer divisions. Sorry I missed your comment until now, it somehow slipped under the (usually reliable) radar :-)

  9. lety4entNo Gravatar Says:

    wow !!!! … really thanks !!! you’re a genius, your algorithm implementation helped me a lot for my homework of Computational Arithmetic in the master degree … :) . Regards

  10. daecNo Gravatar Says:

    Thank you for this program, it is not only interesting, but also helped me quite a bit.

    Have a nice day!

  11. ngnNo Gravatar Says:

    The case for p=2 is wrong again, it should be
    elif p == 2:
    return a

  12. jonNo Gravatar Says:

    I was wondering about the code for p==2 because to return p would return a number that is not mod p.

    Also, would it be better to check if a==0 before calculating the legendre_symbol(a,p)? this function will always return 0 for values of a==0 so that code will never be reached.

  13. jonNo Gravatar Says:

    perhaps on the case p==2 it should be a % p ?

    0**2 mod 2 = 0
    1**2 mod 2 = 1
    2**2 mod 2 = 0
    3**2 mod 2 = 1

    and so on.

Leave a Reply

To post code with preserved formatting, enclose it in `backticks` (even multiple lines)