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 &lt;= n : 
            j = 1
            while(j * j &lt;= 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 &lt;= n * n:
            j = 1
            while (j * j &lt;= n * n):
                if(i * i + j * j == n * n):
                    return True
                j = j + 1
            i = i + 1

        return False

    T = int(input())
    while T&gt;0:
        T-=1
        n = int(input())
        if ((sumSquare(n)) and (Squared(n))):
            print (&quot;Yes&quot;)
        else: 
            print(&quot;No&quot;)
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)) &gt; 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 &gt;= 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.