Finding representation of a given number as a sum of squares

edit

Let us implement the function with bare hands.

def my_reps(k, n):
    if k <= 0 or n == 0:
        return []
    if k == 1:
        return [[isqrt(n)]] if n.is_square() else []
    sols = []
    for xn in [isqrt(n/k)..isqrt(n - k + 1)]:
       shorter_sols = [sol for sol in my_reps(k - 1, n - xn^2) if sol[0] <= xn]  
       sols += [[xn] + s for s in shorter_sols]
    return sols

The cheap recursion makes the code less optimal in cases of a "big $k$". We can check by counting, if the function is doing the right job. We have for instance:

sage: my_reps(3, 1001) [[24, 16, 13], [24, 19, 8], [24, 20, 5], [26, 15, 10], [26, 17, 6], [26, 18, 1], [27, 16, 4], [29, 12, 4], [30, 10, 1], [31, 6, 2]] sage: len(_) 10

Some counting generating function for the solutions is $f^k$, where $f=q+q^4+q^9+q^{16}+q^{25}+\dots$ - but order of components matter. We have collected only decreasing solutions. Fortunately, all displayed solutions have three different components, so we expect a $10\cdot 3!$ count in $f^3$ as coefficient of $q^{1001}$:

sage: R. = PowerSeriesRing(QQ, default_prec=1600, sparse=True) sage: R. = PowerSeriesRing(QQ, default_prec=1600) sage: f = sum([q^(n^2) for n in [1..40]]) + O(q^1601)

And indeed, there is a part with 12*q^1000 + 60*q^1001 + 24*q^1002 + 12*q^1003 + 21*q^1004 + 24*q^1005 + 30*q^1006 + 27*q^1009 in the answer. For $n=1588$ there are particularly few solutions, let us check this degree, too.

sage: my_reps(3, 1588)
[[36, 16, 6]]

An other check may be the brute force loop:

sage: for a, b, c in cartesian_product([[1..40], [1..40], [1..40]]):
....:     if a >= b and b >= c and a^2 + b^2 + c^2 == 1001:
....:         print(a,b,c)
....: 
24 16 13
24 19 8
24 20 5
26 15 10
26 17 6
26 18 1
27 16 4
29 12 4
30 10 1
31 6 2

Again, ten solutions. With a bigger $k$ and a bigger $n$.

sage: sols = my_reps(3, 1234567)
sage: len(sols)
0
sage: sols = my_reps(4, 1234567)
len(sols) 
sage: len(sols)
26234

And after a long time, the list is finally built... The recursive lazy implementation above is of course not optimal. Instead, explicitly implementing an iterator that translates the recursion may be optimal. The number of solution grows quickly with $k$. I do not see any reason for having them stored in memory. Why do you want all those $26234$ solutions for $k=4$ and $n=26234$?

If counting is the purpose (well, avoiding zero, restriction to positive and reject of negatives, restriction to in- or decreasing solutions - all this breaks the arithmetics when counting, but ok...) then a generating series may be the better way.