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’t
September 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