All decompositions of a prime as a sum of four squares

edit

It is straightforward to do it brute-force:

def sum_of_squares(n, number_of_squares):
    if number_of_squares <= 1:
        if n.is_square():
            yield [n.isqrt()]
            if n != 0:
                yield [-n.isqrt()]
        return
    for i in range(n.isqrt(), -1, -1):
        k = n - i**2
        for rest in sum_of_squares(k, number_of_squares - 1):
            yield [i] + rest
            if i != 0:
                yield [-i] + rest

You can check that the theorem holds:

sage: %time len(list(sum_of_squares(389,4)))  
CPU times: user 0.03 s, sys: 0.00 s, total: 0.03 s
Wall time: 0.02 s
3120
sage: 8*(389+1)
3120