## 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:

### 13 Responses to “Computing modular square roots in Python”

1. Steve 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 Ska Says:

Thank you!

3. bob Says:

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

4. eliben Says:

bob,

Fixed, thanks.

5. Seth Says:

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

6. eliben Says:

Seth, I don’t

7. GregS Says:

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

elif p == 2:
return n

with

elif p == 2:
return p
8. eliben 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. lety4ent 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. daec Says:

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

Have a nice day!

11. ngn Says:

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

12. jon 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. jon 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.

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