## Computing modular square roots in Python

March 7th, 2009 at 11:59 amConsider 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:

March 29th, 2009 at 20:50

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) ?

April 13th, 2010 at 12:42

Thank you!

June 2nd, 2011 at 01:49

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

June 3rd, 2011 at 05:25

bob,Fixed, thanks.

July 19th, 2011 at 21:36

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

July 20th, 2011 at 05:36

Seth, I don’tSeptember 19th, 2011 at 15:22

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

with

September 19th, 2011 at 17:18

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

October 9th, 2011 at 17:47

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

November 8th, 2011 at 12:00

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

Have a nice day!

November 2nd, 2012 at 06:51

The case for p=2 is wrong again, it should be

`elif p == 2:`

return a

February 12th, 2014 at 18:41

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.

February 12th, 2014 at 18:47

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.