Sum of two squares in Python

Question

I have written a code based on the two pointer algorithm to find the sum of two squares. My problem is that I run into a memory error when running this code for an input n=55555**2 + 66666**2. I am wondering how to correct this memory error.

def sum_of_two_squares(n):
    look=tuple(range(n))
    i=0
    j = len(look)-1
    while i < j:
        x = (look[i])**2 + (look[j])**2
        if x == n:
            return (j,i)
        elif x < n:
            i += 1
        else:
            j -= 1
    return None

n=55555**2 + 66666**2
print(sum_of_two_squares(n))

The problem Im trying to solve using two pointer algorithm is: return a tuple of two positive integers whose squares add up to n, or return None if the integer n cannot be so expressed as a sum of two squares. The returned tuple must present the larger of its two numbers first. Furthermore, if some integer can be expressed as a sum of two squares in several ways, return the breakdown that maximizes the larger number. For example, the integer 85 allows two such representations 7*7 + 6*6 and 9*9 + 2*2, of which this function must therefore return (9, 2).

Answer 1

You're creating a tuple of size 55555^2 + 66666^2 = 7530713581

So if each element of the tuple takes one byte, the tuple will take up 7.01 GiB.

You'll need to either reduce the size of the tuple, or possibly make each element take up less space by specifying the type of each element: I would suggest looking into Numpy for the latter.

---

Specifically for this problem:

Why use a tuple at all?

You create the variable look which is just a list of integers:

look=tuple(range(n))  # = (0, 1, 2, ..., n-1)

Then you reference it, but never modify it. So: look[i] == i and look[j] == j.

So you're looking up numbers in a list of numbers. Why look them up? Why not just use i in place of look[i] and remove look altogether?

Answer 2

You do not need the ranges at all, and certainly do not need to convert them into tuples. They take a ridiculous amount of space, but you only need their current elements, numbers i and j. Also, as the friendly commenter suggested, you can start with sqrt(n) to improve the performance further.

def sum_of_two_squares(n):
    i = 1
    j = int(n ** (1/2))

    while i < j:
        x = i * i + j * j
        if x == n:
            return j, i

        if x < n:
            i += 1
        else:
            j -= 1

Bear in mind that the problem takes a very long time to be solved. Be patient. And no, NumPy won't help. There is nothing here to vectorize.

Answer 3

As others have pointed out, there's no need to use tuples at all.

One reasonably efficient way of solving this problem is to generate a series of integer square values (0, 1, 4, 9, etc...) and test whether or not subtracting these values from n leaves you with a value that is a perfect square.

You can generate a series of perfect squares efficiently by adding successive odd numbers together: 0 (+1) → 1 (+3) → 4 (+5) → 9 (etc.)

There are also various tricks you can use to test whether or not a number is a perfect square (for example, see the answers to this question), but — in Python, at least — it seems that simply testing the value of int(n**0.5) is faster than iterative methods such as a binary search.

def integer_sqrt(n):
    # If n is a perfect square, return its (integer) square
    # root. Otherwise return -1
    r = int(n**0.5)
    if r * r == n:
        return r
    return -1

def sum_of_two_squares(n):
    # If n can be expressed as the sum of two squared integers,
    # return these integers as a tuple. Otherwise return <None>
    #   i: iterator variable
    #   x: value of i**2
    #   y: value we need to add to x to obtain (i+1)**2
    i, x, y = 0, 0, 1
    # If i**2 > n / 2, then we can stop searching
    max_x = n >> 1
    while x <= max_x:
        r = integer_sqrt(n-x)
        if r >= 0:
            return (i, r)
        i, x, y = i+1, x+y, y+2
    return None

This returns a solution to sum_of_two_squares(55555**2 + 66666**2) in a fraction of a second.