Finding numbers that can be written as the sum of several distinct squared numbers

Question

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

def combs(a):
    if len(a) == 0:
        return [[]]
    cs = []
    for c in combs(a[1:]):
        cs += [c, c+[a[0]]]
    return cs

def main(n):
    k = 0
    l = combs(range(1, int(n**0.5)+1))
    for i in l:
        s = 0
        for j in i:
            s += j**2
        if s == n:
            print(i, n)
            k += 1
    print('\nanswer: ',k)

if __name__ == '__main__':
    main(n = 25)

However, if you replicate my solution, you'll see that after N > 400, the problem gets basically impossible to solve in a short time. So, I was wondering if any optimization exist for this problem?

Answer 1

You can use a standard single-use "coin change" algorithm, with the values of the coins being squares.

def sum_distinct_squares(n):
    W = [1] + [0] * n
    for i in range(1, n+1):
        i2 = i * i
        if i2 > n:
            break
        for j in range(n, i2-1, -1):
            W[j] += W[j - i2]
    return W[n]

print(sum_distinct_squares(100000))

This runs in O(n * sqrt(n)) time, and solves the n=400 case in 0.016s on my machine, and n=100000 in 1.715s.

Answer 2

The following MiniZinc model copes with n=400 (55 solutions) in less than a second:

int: n = 25;
set of int: Domain = 1..ceil(pow(n, 0.5));

%  the array of decision variable decides
%  which integers between 1 and n^0.5 are added as squares
array[Domain] of var bool: b;

constraint n == sum([b[i] * i * i | i in Domain]);

output ["\(n) = "] ++ [if fix(b[i]) then "+\(i)²" else "" endif | i in Domain];

To evaluate the solutions for n=400 by brute force:
Enumerate the 1,048,576 20-bit integers and register those as solutions which yield the desired sum. Each of the twenty bits decides, which integer 1..20 should be squared and added. It does not take that long to loop through a million cases.

Answer 3

Implementation:

from functools import cache
from math import sqrt

@cache
def _square_sums(n, max_i):
    if n == 0:
        return 1
    start = min(max_i, int(sqrt(n)))
    return sum(_square_sums(n - i**2, i - 1) for i in range(start, 0, -1))

def square_sums(n):
    return _square_sums(n, max_i=n)

Tests:

>>> square_sums(25)
2

>>> square_sums(55)
1

>>> square_sums(400)
55

>>> square_sums(10000)
3296089777

>>> square_sums(100000)
2759256389896728737285379

Performance characteristics:

Fast for n < 10000, but slow for much larger n.

Answer 4

Based on the @PaulHankin 's excellent answer, another optimization is also possible using the fast Numba JIT compiler:

import time
import numba
import matplotlib.pyplot as plt

@numba.jit
def sum_distinct_squares_jit(n):
    W = [1] + [0] * n
    for i in range(1, n+1):
        i2 = i * i
        if i2 > n:
            break
        for j in range(n, i2-1, -1):
            W[j] += W[j - i2]
    return W[n]

def sum_distinct_squares_non_jit(n):
    W = [1] + [0] * n
    for i in range(1, n+1):
        i2 = i * i
        if i2 > n:
            break
        for j in range(n, i2-1, -1):
            W[j] += W[j - i2]
    return W[n]

def loop(N, jit = None):
    if jit is True:
        sds = sum_distinct_squares_jit
    if jit is False:
        sds = sum_distinct_squares_non_jit
  
    times = []
    t_i = time.clock()
    for number in range(1, N):
        sds(number)
        times.append(time.clock() - t_i)
    plt.plot(times)

loop(N = 10_000, jit = True)
loop(N = 10_000, jit = False)
plt.legend(['JIT', 'None JIT'])
plt.show()

Running this in a for loop and plotting the time it takes for JIT and Non-JIT, we get:

So, using JIT is significantly faster. However, this comes with a price: Numba doesn't support the bigint library and for big numbers using Numba is not valid.

Answer 5

There is only a finite set of numbers not sum of distinct squares (http://oeis.org/A001422)

So here is an O(1) solution :

If N not in [2, 3, 6, 7, 8, 11, 12, 15, 18, 19, 22, 23, 24, 27, 28, 31, 32, 33, 43, 44, 47, 48, 60, 67, 72, 76, 92, 96, 108, 112, 128]:
  Print("N is a sum of distinct squares")