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.

• 200_success: Welcome to Code Review! I've attempted to clarify the problem. If I've misinterpreted it, please correct it.
• Clockwork-Muse: I'd probably start with something like total - ((int) sq_root(total)) ^ 2 - ie, subtract the largest square available. I'm not sure if this always gets you the fewest number of elements, though. Your big problem is that at every step, you keep trying 1, or other numbers too small for the remaining number of 'attempts'.
• John Palmer: Switching from seq.* to array.* will give a pretty significant perf boost.

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'

• Andy B: Thanks, this works great and is really quick. Thanks also for the really detailed explanation. I also came up with another solution using the sqrt as suggested by @Clockwork-Muse which I have put below.
• Andy B: Actually I deleted the answer because for a number such as 48 it was taking 5 numbers to calculate it and didn't take into account the maximum number of values to use