Tonelli shanks algorithm is used for finding modular square root of number n under mod p in minimal time
def is_quadratic_residue(n,p):
'''check if the n is qudratic residue in mod p using legendre symbol'''
if n%p==0:
return True
return pow(n,(p-1)//2,p)==1
def tonelli_shanks(n,p):
'''find the modular square root of n in mod p if possible'''
if n%p==0:
# when n = o(mod p) square root of 0 is 0
return 0
if not is_quadratic_residue(n,p):
print("square root does not exits")
return None
if p%3==4:
return pow(n,(p+1)//4,p)
# for all p%4==1 we will apply tonelli shanks algorithm
# step 1
# p-1 = 2^s.q
q=p-1
s=0
while(q%2==0):
s+=1
q=q//2
# step 2
n_res = 2
while is_quadratic_residue(n_res,p):
n_res+=1
print(n_res)
# step 3
m=s
c=pow(n_res,q,p)
t=pow(n,q,p)
r=pow(n,(q+1)//2,p)
# step 4
while(t!=1):
i=0
temp=t
while(temp!=1):
i+=1
temp = (temp*temp)%p
pow2 = pow(2,m-i-1)
b=pow(c,pow2,p)
m=i
c = (b**2)%p
t = (t*b**2)%p
r = (r*b)%p
return r
x=tonelli_shanks(5,41)
print(x)