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?

  • \$\begingroup\$ Could you clarify what you're asking, please? \$\endgroup\$ – jonrsharpe Dec 5 '14 at 16:20
  • \$\begingroup\$ Just out of curiosity, does this have any real world application? \$\endgroup\$ – Bruno Costa Dec 5 '14 at 16:23
  • \$\begingroup\$ @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? \$\endgroup\$ – john mangual Dec 5 '14 at 16:23
  • 1
    \$\begingroup\$ @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. \$\endgroup\$ – john mangual Dec 5 '14 at 16:30
  • \$\begingroup\$ You may want to first make sure your number is not in sequence A004215 in OEIS, see here for details. \$\endgroup\$ – 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)

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

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.

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.