Recently, I came across a really interesting question:

Given the number N, how many combinations exist that can be written as the sum of several distinct squared numbers?

For example, 25 can be written as:

25 = [5**2 , 3**2   4**2]

So the answer would be 2.

But for 31, there exist no solutions, so the answer would be 0.

At the beginning, I thought this would be an easy problem to solve. I just assumed that if I count every combination of squared numbers less than square root of number N, then I should be fine, time-wise. (I found the combinations based on

Explanation:

• Work backwards from square candidates (e.g. start from i=5 for n=25).
  • If i**2 == n, yield a result.
  • If i**2 < n - i**2, stop.
  • Otherwise, yield results from valid square sums of n - i**2.

Implementation:

from math import sqrt

def square_sums(n):
    if n == 0:
        yield []
        return
    start = min(n - 1, int(sqrt(n)))
    for i in range(start, 0, -1):
        if i**2 < n - i**2:  # Enforces ordering.
            return
        results = (xs   [i] for xs in square_sums(n - i**2))
        if i**2 == n - i**2:  # Disallow duplicates by dropping first result.
            try:
                next(results)
            except StopIteration:
                pass
        yield from results

Tests:

>>> list(square_sums(25))
[[5], [3, 4]]

>>> list(square_sums(400))
[[20], [12, 16]]

>>> len(list(square_sums(10000)))
23015