# Listing the representations of integer M as a sum of three squares M = x^2 + y^2 + z^2

I am trying a number as the sum of 3 squares. For my particular project I will require a complete list since the next step is to compute nearest neighbors (here).

In any case we'd like to find $M = x^2 + y^2 + z^2$ in all possible ways. Here's what I wrote in Python:

M = 1062437
m = int(np.sqrt(M))

x = np.arange(int(np.sqrt(m)))
y = M - x**2

sq = [(i,j, int(np.sqrt(M - i**2 - j**2))) for i in range(m) for j in range(m) if i**2 + j**2 in y]


This is easy to explain in pseudocode.

• build a list of remainders $\{ M - k^2 : 0 \leq k \leq \sqrt{M} \}$
• loop over all possible $0 < i,j < \sqrt{M}$ check that $i^2 + j^2$ is in the list

This is one of several possible strategies. In my case, I don't even need a fixed M it could be a range, such as 10**6 < M < 10**6 + 10*4.

I do not mind looping over values of M at this point. Can we recycle our computations somehow?

• Could you clarify what you're asking, please? – jonrsharpe Dec 5 '14 at 16:20
• Just out of curiosity, does this have any real world application? – Bruno Costa Dec 5 '14 at 16:23
• @jonrsharpe 1) is this good python and algorithms? 2) if I let M run over a range if there a way to avoid computing i^2 + j^2 over and over? – john mangual Dec 5 '14 at 16:23
• @BrunoCosta these might appear in cryptography or coding theory. Point sets and nearest neighbor algorithms appear all over. machine learning, digitals maps, front-end, back-end. – john mangual Dec 5 '14 at 16:30
• You may want to first make sure your number is not in sequence A004215 in OEIS, see here for details. – Jaime Dec 5 '14 at 17:24

You might be interested in this question.

While looking for solution to

$$x^{2} + y^{2} + z^{2} = M$$

you can assume

$$x \leq y \leq z$$.

and perform rotation if needed afterward.

This is good because, you can start iterating over x in a smaller range. You know that

$$M = x^{2} + y^{2} + z^{2} \ge x^{2} + x^{2} + x^{2} = 3 * x^{2}$$

, therefore $$x \le sqrt(M / 3)$$

Similarly, when iterating over y, you can start at xand go up to a smaller higer bound because we have :

$$M = x^{2} + y^{2} + z^{2} \ge x^{2} + y^{2} + y^{2} = x^{2} + 2 * y^{2}$$

, therefore $$x \le y \le sqrt((M - x^{2})/ 2)$$

Then, when once x and y are fixed, you don't need to iterate over z : just compute

$$z = sqrt(M - x^{2} - y^{2})$$

and check that it is indeed an integer.

Here's a proof-of-concept piece of code : I've altered your code a bit to be able to compare the results you originally had with the one I get :

sq = set(tuple(sorted([i,j, int(np.sqrt(M - i**2 - j**2))])) for i in range(m) for j in range(m) if i**2 + j**2 in y)
print(len(sq))

sol = []
lim_x = int(np.sqrt(M/3))
for x in range(1 + lim_x):
rem_x = M - x*x
lim_y = int(np.sqrt(rem_x / 2))
for y in range(x, 1 + lim_y):
rem_y = rem_x - y*y
z = int(np.sqrt(rem_y))
if z*z == rem_y:
assert x <= y <= z
assert x*x + y*y + z*z == M
sol.append((x, y, z))
print(len(sol))

assert all(s in sol for s in sq)


You have 26 different ordered solutions out of the 290 I found (in a smaller time range).

Finally, you might be able to adapt the equations above (and the code written from it) if you want to handle ranges of M values.

I don't have np (I assume numpy) installed. So I tried pure Python. I tried to store i,j,i**2+j**2 into a list of tuples and then reuse it later. That made it slower than your code. I also tried (i,j) as dict key and i**2+j**2 as value. That did not help either.

The only thing that made it faster was to have y as a set.

m = int(math.sqrt(M))
x = range(int(math.sqrt(m)))
y = set([M - x**2 for x in x])
return [(i,j, int(math.sqrt(M - i**2 - j**2))) for i in range(m) for j in range(m) if i**2 + j**2 in y]

• Code uses variables x, i and j instead of the x, y and z from the problem description. It would be clearer to use the same.
• Calculating sqrt(m) looks like a bug because m = sqrt(M) already.
• You could avoid computing the squares over and over again in the list comprehension by keeping x**2 in a variable and making i and j loop over it.
• You could use for j in range(i, m) to avoid checking the same pair of numbers in reverse. If you need all the permutations of each solution, generate the reversed version at a later stage.
• Checking in y on an array requires a linear search. Even though the search is implemented in fast C code, a Python set is faster when the array is large enough.