# Determine whether a number and its square are both the sum of two squares

Code gets timed out and needs to be simplified.

To ensure that n = sum of squares , To determine if n^2 = sum of squares , I need to input 4 integers : 5881, 2048, 2670, 3482

#To define the sum of squares
def sumSquare(n) :
i = 1
while i * i <= n :
j = 1
while(j * j <= n) :
if (i * i + j * j == n) :
return True
j = j + 1
i = i + 1

return False

#To define that the square of n is equals to the square of i added to the square of j
def Squared(n):
i = 1
while i * i <= n * n:
j = 1
while (j * j <= n * n):
if(i * i + j * j == n * n):
return True
j = j + 1
i = i + 1

return False

T = int(input())
while T>0:
T-=1
n = int(input())
if ((sumSquare(n)) and (Squared(n))):
print ("Yes")
else:
print("No")

• $n^2 = n^2 + 0^2$ is always a sum of two squares... – Stefan Pochmann Jan 21 at 23:40
• You may also improve the performance by using formula from this post: math.stackexchange.com/a/1461742 – tsh Jan 22 at 9:52
• Welcome to Code Review! Does your code work for smaller numbers? – Malachi Jan 22 at 17:19
• it works faster for small numbers as the initial i is 1. then it keeps adding 1 until all conditions are met – Jiwoo Kang Jan 22 at 17:47

## 2 Answers

Style

Various tiny improvement can be made to make your code more aesthetically pleasing without changing the actual behavior:

• removing trailing whitespace

• removing useless parenthesis

• removing whitespace before colon

• improving the function names for something which brings more meaning

• defining docstrings for your methods

I would highly recommend having a look the the Style Guide for Python code called https://www.python.org/dev/peps/pep-0008/ .

At this stage, I have:

def is_sum_of_two_squares(n):
"""Check whether n can be written as i*i + j*j."""
i = 1
while i * i <= n:
j = 1
while j * j <= n:
if i * i + j * j == n:
return True
j = j + 1
i = i + 1
return False

def n_squared_is_sum_of_two_squares(n):
"""Check whether n*n can be written as i*i + j*j."""
i = 1
while i * i <= n * n:
j = 1
while j * j <= n * n:
if i * i + j * j == n * n:
return True
j = j + 1
i = i + 1
return False

for n in [5881, 2048, 2670, 3482]:
print(n, is_sum_of_two_squares(n), n_squared_is_sum_of_two_squares(n))


Deleting code

One can see that n_squared_is_sum_of_two_squares could be written by reusing is_sum_of_two_squares:

def n_squared_is_sum_of_two_squares(n):
"""Check whether n*n can be written as i*i + j*j."""
return is_sum_of_two_squares(n*n)


(Maybe we could get rid of the function altogether)

Now, we have a clearer view on the code to be optimised.

Optimisations

Knowing that we are looking for i such that i * i < n is a good idea. Unfortunately, we do it in a slow way as compute additions and multiplications at each step.

We could compute the square root once and be done:

import math

def is_sum_of_two_squares(n):
"""Check whether n can be written as i*i + j*j."""
s = int(math.sqrt(n))
for i in range(1, 1 + s):
j = 1
while j * j <= n:
if i * i + j * j == n:
return True
j = j + 1
return False


Another observation is that once you are testing for a given value of i, there is not much option for the possible values of j. There is only one possible value:

def is_sum_of_two_squares(n):
"""Check whether n can be written as i*i + j*j."""
s = int(math.sqrt(n))
for i in range(1, 1 + s):
j2 = n - i * i
j = int(math.sqrt(j2))
if j * j == j2 and j > 0:
return True
return False


Which brings a huge optimisation.

A final optimisation can be to reduce once more the search space for i. A common technique can be to assume that we are looking for (i, j) such that i²+j²=n and i <= j. This works because the problem is symmetrical and any solution (i, j) leads to a solution (j, i).

Then, we can stop at $$\\sqrt{\frac{n}{2}}\$$ instead of $$\\sqrt{n}\$$.

def is_sum_of_two_squares(n):
"""Check whether n can be written as i*i + j*j."""
# Assume 1 <= i <= j
# Then n = i*i + j*j >= 2*i*i
# So i <= sqrt(n/2)
s = int(math.sqrt(n/2))
for i in range(1, 1 + s):
j2 = n - i * i
j = int(math.sqrt(j2))
if j * j == j2 and j > 0:
return True
return False


Based on comments from AJNeufeld, one can actually go further with this idea, iterating over the "j-range" which is actually smaller:

def is_sum_of_two_squares2(n):
"""Check whether n can be written as i*i + j*j."""
# Assume 1 <= i <= j
# Then n = i*i + j*j becomes
#      n <= 2*j*j which implies j >= sqrt(n/2)
# And  n >= 1*1 + j*j which implies j <= sqrt(n-1)
for j in range(int(math.sqrt(n/2)), 1 + int(math.sqrt(n-1))):
i2 = n - j * j
i = int(math.sqrt(i2))
if i * i == i2 and i > 0:
return True
return False


Also, yet another brilliant idea from AJNeufeld is to use math.isqrt:

def is_sum_of_two_squares(n):
"""Check whether n can be written as i*i + j*j."""
# Assume 1 <= i <= j
# Then n = i*i + j*j becomes
#      n <= 2*j*j which implies j >= sqrt(n/2)
# And  n >= 1*1 + j*j which implies j <= sqrt(n-1)
for j in range(math.isqrt(n//2), 1 + math.isqrt(n-1)):
i2 = n - j * j
i = math.isqrt(i2)
if i * i == i2 and i > 0:
return True
return False

• s = math.isqrt(n / 2) and j = math.isqrt(j2) – AJNeufeld Jan 21 at 15:16
• You're going from the wrong end. for i in range(1, 1+s): means trying $\sqrt{n/2}$ cases, with j barely changing between each case. If you varied j from $\sqrt{n/2}$ to $\sqrt{n}$, you'll only need to try $\sqrt{n} - \sqrt{n/2}$ cases, with the computed values of i varying significantly at each step. As an example, checking $5881^2$ requires up to 4158 iterations when looping over i, but only 1722 when looping over j. – AJNeufeld Jan 21 at 16:00
• @AJNeufeld Great idea, I'll give it a try when I have a bit of time! – SylvainD Jan 21 at 19:31

You can speed it up by saving the squares in a set. If you wrote a function to do this, you can give it a mutable default argument for it to save all the squares instead of having to calculate them repeatedly.

# dict used because there’s no insertion ordered set
def sum_squares(n, saved_squares=dict()):
# check for sum while saving squares for future calls
check_end = n
if int(n**(1/2)) > len(saved_squares):
check_end = (len(saved_squares) + 1)**2
found_while_saving = False
for num_to_save in range(len(saved_squares) + 1, int(n**(1/2) + 1)):
square = num_to_save**2
saved_squares[square] = None
if n - square in saved_squares:
found_while_saving = True
if found_while_saving == True:
return True
# check for sum in set of already calculated squares
for square in saved_squares:
if n - square in saved_squares:
return True
# insertion order needed so early return won’t happen
if square >= check_end:
return False


With the numbers 5881, 2048, 2670, and 3482, it will only have to calculate any squares for the first call with the number 5881. On the other three calls it can skip to checking the set of all the already calculated squares. But no matter what order of numbers it’s called with, it will only have to calculate the square of any number once.

This is how I tested it on repl.it:

nums_to_check = (2048, 2670, 3482, 5881)
for n in nums_to_check:
print(sum_squares(n) and sum_squares(n**2))


It completed in under a second, so hopefully it’s fast enough for the site you’re using.

• Fine; but also mutable default arguments are considered risky and some linters will point that out to you. The problem is that if the default object actually is mutated, that will surprise callers e.g. from other modules that expect a certain default (in this case an empty dictionary). The workaround is to accept either None or a sentinel object, and if that's seen, make the dict. – Reinderien Jan 21 at 3:12
• Details about Multable default arguments @Reinderien mentioned. For reference to future readers. – Ch3steR Jan 21 at 11:28
• Thanks. I think it should be alright to dismiss any warning about it for this since it’s only being used as a way to store values for future calls. I’ve read that’s one of the things it’s intended for. But maybe that’s not true, so correct me if I’m wrong about that. – my_first_c_program Jan 22 at 1:02