Decompose a value as a sum of square numbers

Question

My daughter had a question on her maths homework which was to write a value as a sum of 4 or fewer square numbers.

For example, x = 73 could be 36 + 36 + 1.

I came up with a really brute force algorithm that shows all combinations but when the values get above 1000 it becomes quite slow.

Is there a clever way of achieving this? Here is my algo in F#

let squares = [ 0.. 100 ] |> Seq.map (fun x -> x * x) |> Seq.toList

let rec calc total attempts squares accu results =
  match (total, attempts, squares) with
  | (0, _, _) -> accu :: results
  | (_, 0, _) -> [] :: results
  | (x, _, _) when x < 0 -> [] :: results
  | (_, _, []) -> [] :: results
  | total, attempts, squares ->
    let filteredSquares = squares |> Seq.filter ((>=) total )
    filteredSquares 
      |> Seq.collect(fun sq -> 
                          calc (total - sq)
                               (attempts - 1)
                               (filteredSquares |> Seq.toList)
                               (sq :: accu)
                               results) |> Seq.toList

let res = 
  calc 8058 4 squares [] [] 
  |> Seq.filter(fun lst -> lst <> [])
  |> Seq.sortBy (fun lst -> lst.Length)
  |> Seq.take 1
  |> Seq.toList 

Additional comments on better F# code would be appreciated as well. One thing I was wondering was whether I could make it lazy using sequences.

Answer 1

let squares = [ 0.. 100 ] |> Seq.map (fun x -> x * x) |> Seq.toList

I don't think limiting the squares this way is a good idea. Instead, if you moved it inside calc, you could write it as:

let squares =
    Seq.initInfinite id
    |> Seq.map (fun x -> x * x)
    |> Seq.takeWhile (fun x -> x <= total)
    |> Seq.toList

---

calc is a pretty bad name: it doesn't say anything about what the function calculates.

---

If you limit yourself to returning always only the shortest sum, the signature of the method could be simplified quite a lot: you could get rid of attempts, accu and results. This would also somewhat simplify the implementation.

And even if you don't do that, I think that filtering out empty lists should be done inside calc, not outside.

---

If you simplify the signature, you could then apply memoization to make the computation much faster. Memoization will help here, because what's slowing down the computation is calculating the result for the same total over and over.

The implementation could look like this:

let rec calculateSquaresSum =
    let calculateSquaresSum' total =
        if total = 0 then
            []
        else
            let squares =
                Seq.initInfinite (fun x -> x + 1)
                |> Seq.map (fun x -> x * x)
                |> Seq.takeWhile (fun x -> x <= total)
                |> Seq.toList
                |> List.rev

            squares
            |> Seq.map (fun sq -> sq :: (calculateSquaresSum (total - sq)))
            |> Seq.minBy (fun sum -> sum.Length)

    memoize calculateSquaresSum'

This will calculate the result for values around 10.000 immediately, but 100.000 takes a few seconds, so there likely is a more efficient way to implement this.

---

Though if you're okay with finding just some decomposition with at most the given length, not necessarily the shortest one, having the attempts parameter makes sense and makes the computation much faster when the decomposition is possible. (E.g. calculateSquaresSum(1000000 - 1, 4) returns immediately with [998001; 1849; 100; 49], while calculateSquaresSum(100000 - 1, 3) takes few seconds to return with a negative answer.)

The code:

let rec calculateSquaresSum =
    let calculateSquaresSum' (total, attempts) =
        if total = 0 then
            Some []
        elif attempts = 0 then
            None
        else
            let squares =
                Seq.initInfinite (fun x -> x + 1)
                |> Seq.map (fun x -> x * x)
                |> Seq.takeWhile (fun x -> x <= total)
                |> Seq.toList
                |> List.rev

            squares
            |> Seq.map (fun sq ->
                match calculateSquaresSum (total - sq, attempts - 1) with
                | None -> None
                | Some tail -> Some (sq::tail))
            |> Seq.tryPick id

    memoize calculateSquaresSum'