# Square root a number

The task is as follows: Given a positive integer x and a prime n > x, output the smallest positive integer y such that (y * y) mod n = x. An important part of this question is the time limit specified below which excludes brute force solutions.

If there is no such value y then your code should output N.

Test cases

(2, 5, N),
(3, 5, N),
(4, 5, 2),
(524291, 1048583, N),
(529533, 1048583, N),
(534775, 1048583, 436853),
(540017, 1048583, 73675),
(536870913, 1073741827, 375394238),
(542239622, 1073741827, 267746399),
(547608331, 1073741827, N),
(552977040, 1073741827, 104595351),
(1099511627676, 1099511627791, N),
(1099511627677, 1099511627791, 269691261521),
(1099511627678, 1099511627791, 413834069585),
(1267650600228229401496703204376, 1267650600228229401496703205653, 5312823546347991512233563776),
(1267650600228229401496703204476, 1267650600228229401496703205653, N)
(1267650600228229401496703204576, 1267650600228229401496703205653, N)
(1267650600228229401496703204676, 1267650600228229401496703205653, 79905476259539917812907012931)


Input and output

You may take input and give output in any way that is convenient. If you don't like to output N then any Falsey value will do.

Restrictions

Your code must answer all the test cases in under 1 minute on a standard desktop. This time limit is only to prevent brute force answers and I expect good answers to run almost instantly. You may not use any library or builtin that solves this problem or that tests if a number is a quadratic residue.

• So languages without support for big-integer data type are ruled out. Pity Commented Jul 7, 2017 at 14:12
• @LuisMendo Not if you can code up support for 1267650600228229401496703205653 yourself or if you have 128 bit support such as __int128 in gcc. There are also a number of 256 bit int libraries out there for various languages. Last, lots of languages have an arbitrary precision int library.
– user9206
Commented Jul 7, 2017 at 14:13

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

• %timeit Q(1267650600228229401496703204676,1267650600228229401496703205653) 100 loops, best of 3: 2.83 ms per loop :)
– user9206
Commented Jul 7, 2017 at 19:12
• What a great answer! You have restored my faith in PPCG.
– user9206
Commented Jul 7, 2017 at 19:14
• Excuse the newbie question, but whats PPCG? Polish Python Coders Group? Commented Jul 10, 2017 at 8:49
• @DaveBoltman Programming Puzzles & Code Golf.
– orlp
Commented Jul 10, 2017 at 9:26

# Pyth, 83 82 bytes

=eAQM.^GHQKf%=/H=2;1=gftgT/Q;1HJg~gGHh/H2WtG=*J=gT^2t-K=Kfq1gG^2T1=%*G=^T2Q;hS%_BJ


Test suite

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

• That reduction is quite funny :)
– user9206
Commented Jul 10, 2017 at 17:09

Usually I like brute force answers.. Since this is strongly discouraged by the time limit, here's the most efficient way I know of:

r p=((\r->min(mod(-r)p)r)$).f p f p x|p==2=x|q x=f 0 0|mod p 4==3=x&div(p+1)4|let(r,s)=foldl i(p-1,0)[1..t 1o]=m$x&(d$r+1)*(b&d s)where q a=1/=a&o;m=(modp);a&0=1;a&e|even e=m$a&d e^2|0<1=m\$(a&(e-1))*a;b=[n|n<-[2..],q n]!!0;i(a,b)_|m(x&d a*b&d b)==p-1=(d a,d b+o)|0<1=(d a,d b);o=d p;t y x|even x=t(y+1)(d x)|0<1=y;d=(div2)


Try it online!

I'm sure this can be golfed down further, but this should do for now.

• Could you modify the TIO code so it gives the answers as output? I just get "True" at the moment.
– user9206
Commented Jul 10, 2017 at 8:22
• Try it online! Commented Jul 10, 2017 at 9:31
• @Lembik You need to replace testCases with those from the original TIO, it barely fit into a comment even without them. Commented Jul 10, 2017 at 9:34
• @ØrjanJohansen Thanks a lot! I adjusted my answer with your code and replaced testCases. Commented Jul 10, 2017 at 15:11
• Huh, I'm seeing a bizarre bug with that TIO link - if I click it it has the code but neither running nor getting the URL from the menu options works - but if I copy the URL from the address bar and paste it into a different tab, then it works. Commented Jul 10, 2017 at 21:29