Square root a number
Python 2, 166 bytes
def Q(x,n,a=0):
e=n/2
while pow(a*a-x,e,n)<2:a+=1
w=a*a-x;b=r=a;c=s=1
while e:
if e%2:r,s=(r*b+s*c*w)%n,r*c+s*b
b,c=(b*b+c*c*w)%n,2*b*c;e/=2
return min(r,-r%n)
Pyth, 83 82 bytes
=eAQM.^GHQKf%=/H=2;1=gftgT/Q;1HJg~gGHh/H2WtG=*J=gT^2t-K=Kfq1gG^2T1=%*G=^T2Q;hS%_BJ
This program implements the Tonelli-Shanks algorithm. I wrote it by closely following the Wikipedia page. It takes as input (n, p)
.
The absence of a square root is reported by the following error:
TypeError: pow() 3rd argument not allowed unless all arguments are integers
This is very intricately golfed code, written in the imperative style, as opposed to the more common functional style of Pyth.
The one subtle aspect of Pyth I'm using is =
, which, if not immediately followed by a variable, searches forward in the program for the next variable, then assigns the result of the following expression to that variable, then returns that result. I will refer throughout the explanation to the wikipedia page: Tonelli-Shanks algorithm, as that is the algorithm I am implementing.
Explanation:
=eAQ
A
takes a 2-tuple as input, and assigns the values to G
and H
respectively, and returns its input. Q
is the initial input. e
returns the last element of a sequence. After this snippet, G
is n
, and H
and Q
are p
.
M.^GHQ
M
defines a 2 input function g
, where the inputs are G
and H
. .^
is Pyth's fast modular exponentiation function. This snippet defines g
to mean exponentiation mod Q
.
Kf%=/H=2;1
f
defines a repeat until false loop, and returns the number of iterations it runs for, given 1
as its input. During each iteration of the loop, we divide H
by 2, set H
to that value, and check whether the result is odd. Once it is, we stop. K
stores the number of iterations this took.
One very tricky thing is the =2;
bit. =
searches ahead for the next variable, which is T
, so T
is set to 2. However, T
inside an f
loop is the iteration counter, so we use ;
to get the value of T
from the global environment. This is done to save a couple of bytes of whitespace that would otherwise be needed to separate the numbers.
After this snippet, K
is S
from the wikipedia article (wiki), and H
is Q
from the wiki, and T
is 2
.
=gftgT/Q;1H
Now, we need to find a quadratic nonresidue mod p
. We'll brute force this using the Euler criterion. /Q2
is (p-1)/2
, since /
is floored division, so ftgT/Q;1
finds the first integer T
where T ^ ((p-1)/2) != 1
, as desired. Recall that ;
again pulls T
from the global environment, which is still 2. This result is z
from the wiki.
Next, to create c
from the wiki, we need z^Q
, so we wrap the above in g ... H
and assign the result to T
. Now T
is c
from the wiki.
Jg~gGHh/H2
Let's separate out this: ~gGH
. ~
is like =
, but returns the original value of the variable, not its new value. Thus, it returns G
, which is n
from the wiki.
This assigns J
the value of n^((Q+1)/2)
, which is R
from the wiki.
Now, the following takes effect:
~gGH
This assigns G
the value n^Q
, which is t
from the wiki.
Now, we have our loop variables set up. M, c, t, R
from the wiki are K, T, G, J
.
The body of the loop is complicated, so I'm going to present it with the whitespace, the way I wrote it:
WtG
=*J
=
gT^2
t-
K
=Kfq1gG^2T1
=%*G=^T2Q;
First, we check whether G
is 1. If so, we exit the loop.
The next code that runs is:
=Kfq1gG^2T1
Here, we search for the first value of i
such that G^(2^i) mod Q = 1
, starting at 1. The result is saved in K
.
=gT^2t-K=Kfq1gG^2T1
Here, we take the old value of K
, subtract the new value of K
, subtract 1, raise 2 to that power, and then raise T
to that power mod Q
, and then assign the result to T
. This makes T
equal to b
from the wiki.
This is also the line which terminates the loop and fails if there is no solution, because in that case the new value of K
will equal the old value of K
, 2 will be raised to the -1
, and the modular exponentiation will raise an error.
=*J
Next, we multiply J
by the above result and store it back in J
, keeping R
updated.
=^T2
Then we square T
and store the result back in T
, setting T
back to c
from the wiki.
=%*G=^T2Q
Then we multiply G
by that result, take it mod Q
and store the result back in G
.
;
And we terminate the loop.
After the loop's over, J
is a square root of n
mod p
. To find the smallest one, we use the following code:
hS%_BJ
_BJ
creates the list of J
and its negation, %
implicitly takes Q
as its second argument, and uses the default behavior of Pyth to apply % ... Q
to each member of the sequence. Then S
sorts the list and h
takes its first member, the minimum.
SageMath, 93 bytes
By reduction to a much harder problem for which SageMath happens to have fast enough builtins.
def f(x,n):g=primitive_root(n);k=Zmod(n)(x).log(g);y=pow(g,k//2,n);return k%2<1and min(y,n-y)
Try it on SageMathCell