# return the sum of the squares that is equal to n²

I'm struggling to solve this Kata. My code works but I can't pass because it's timing out. I have googled it and even tried other people's code but still doesn't pass.

Any suggestion on how can I make a faster solution?

My little sister came back home from school with the following task: given a squared sheet of paper she has to cut it in pieces which, when assembled, give squares the sides of which form an increasing sequence of numbers. At the beginning it was lot of fun but little by little we were tired of seeing the pile of torn paper. So we decided to write a program that could help us and protects trees.

Given a positive integral number n, return a strictly increasing sequence (list/array/string depending on the language) of numbers, so that the sum of the squares is equal to n².

If there are multiple solutions (and there will be), return as far as possible the result with the largest possible values:

Examples

decompose(11) must return [1,2,4,10]. Note that there are actually two ways to decompose 11²: 11² = 121 = 1 + 4 + 16 + 100 = 1² + 2² + 4² + 10² but don't return [2,6,9], since 9 is smaller than 10.

For decompose(50) don't return [1, 1, 4, 9, 49] but [1, 3, 5, 8, 49] since [1, 1, 4, 9, 49] doesn't form a strictly increasing sequence.

Note: Neither [n] nor [1,1,1,…,1] are valid solutions. If no valid solution exists, return nil, null, Nothing, None (depending on the language) or "[]" (C) ,{} (C++), [] (Swift, Go).

The function "decompose" will take a positive integer n and return the decomposition of N = n² as:

• [x1 ... xk] or
• "x1 ... xk" or
• Just [x1 ... xk] or
• Some [x1 ... xk] or
• {x1 ... xk} or
• "[x1,x2, ... ,xk]"

depending on the language (see "Sample tests")

decompose 50 returns "1,3,5,8,49"
decompose 4 returns "Nothing"

Hint: Very often xk will be n-1.

My code:

def decompose(n):
count = 1
lst = []
tot = n ** 2

while tot > 0:
if n - count == 0:
return None
elif (n - count) ** 2 <= tot:
tot -= (n - count) ** 2
lst.append(n - count)
count += 1
else:
count += 1

return sorted(lst)

• I doubt it works. Better try stackoverflow.com. For every count you have to add either (n - count) or skip it. In the first case you could recursively call to see whether there is a solution, and then that is the best solution. Also count is not limited, for every else. Jun 17, 2021 at 1:18
• I have posted on StackOverflow and other users told me to post it here. The code works, I have tested with different values, the problem is when it is tested with big numbers, and it takes a long time to run (starts to get slow with 7 digits). The way that I´m using the 'while' isn´t in a recursively? I´m new to coding and I still get confused with some terms. I appreciate your help.
– RMEP
Jun 17, 2021 at 1:54
• – vnp
Jun 17, 2021 at 4:50
• I've just tried the posted code, and decompose(50) returns None. What testing have you done? Jun 17, 2021 at 7:10
• @vnp the theorem you linked isn't applicable because of the side conditions in the problem. E.g. strictly increasing sequence (same number can't occur twice) and select the solution that has the largest possible numbers when multiple solutions exist. The linked theorem says you can write all natural numbers as the sum of four squares, not that it's not possible to have more or fewer squares too. There are in fact several solutions. Jun 17, 2021 at 10:32

As Joop's answer says, it makes sense to do this kind of search using recursion. That gives a natural means of backtracking when one line of search doesn't work out.

If we consider candidates in order from largest to smallest, we naturally know the correct order - just prepend each successful element to the list, as it will be smaller than all the previously-added ones.

I split the problem into two functions. The first is a recursive function to find the distinct squares no larger than c which sum to n. The second is a function which calls that with the arguments n² and n - 1:

def helper(n, c):
"""
Find a set of distinct numbers, each no greater than C, whose
squares sum to N.

Returns as a list, in numerical order.

If no such set exists, return None.
"""

def decompose(n):
return helper(n * n, n - 1)


Then to fill in the body of helper, I did the recursive step first. Assume that c is in the list; then we need to find the numbers smaller than c whose squares sum to n - c²:

    return helper(n - c * c, c - 1) + [c]


If that doesn't find a solution, then the recursive call will return None and we can detect that by catching the exception that None + [c] will throw:

    try:
return helper(n - c * c, c - 1) + [c]
except TypeError:
# the recursive call returned None, so try the next smaller
return helper(n, c - 1)


Now we need to ensure that the recursion terminates:

    if n == 0:
# found an exact sum
return []

    if c == 0:
# ran out of candidates
return None


And ensure that we don't waste time exploring solutions where n - c² goes negative:

    if n < c * c:
c = int(math.sqrt(n))


Putting it all together:

import math

def helper(n, c):
"""
Find a set of distinct numbers, each no greater than C, whose
squares sum to N.

Returns as a list, in numerical order.

If no such set exists, return None.
"""
if n == 0:
return []
if n < c * c:
c = int(math.sqrt(n))
if c == 0:
return None
try:
return helper(n - c * c, c - 1) + [c]
except TypeError:
# the recursive call returned None, so try the next smaller
return helper(n, c - 1)

def decompose(n):
return helper(n * n, n - 1)

if __name__ == "__main__":
print(decompose(11))
print(decompose(50))

• Great answer! Did not know you dabbled with the snake ;-) Jun 17, 2021 at 10:25
• @N3buchadnezzar just a dabbler, I'm afraid. But thank you! :-) Jun 17, 2021 at 16:29

Just have a java at hand.

List<Integer> decompose(int n) {
int count = 1;
List<Integer> lst = new ArrayList<>();
int tot = n * n;
while (tot > 0) {
if (n - count == 0) {
return Collections.emptyList();
} else if ((n - count) * (n - count) <= tot) {
tot -= (n - count) * (n - count);
count++;
} else {
count++;
}
}
Collections.sort(lst);
return lst;
}


My function as I think you intended:

List<Integer> decompose2(int n) {
return decompose2Rec(n - 1, n * n);
}

List<Integer> decompose2Rec(int i, int tot) {
if (tot == 0) {
return new ArrayList<>();
}
if (i <= 0 || tot < 0) {
return null;
}
List<Integer> sublist = decompose2Rec(i - 1, tot - i*i);
if (sublist != null) {
return sublist;
}
return decompose2Rec(i - 1, tot);
}


Test code:

    for (int n = 11; n < 15; ++n) {
System.out.printf("%3d -> %s%n    -> %s%n", n, decompose(n), decompose2(n));
}


Result, the first line your function:

 11 -> [1, 2, 4, 10]
-> [1, 2, 4, 10]
12 -> []
-> [1, 2, 3, 7, 9]
13 -> [5, 12]
-> [5, 12]
14 -> []
-> [4, 6, 12]

150 -> [1, 3, 17, 149]
-> [1, 3, 17, 149]
151 -> []
-> [3, 6, 16, 150]
152 -> [1, 2, 3, 17, 151]
-> [1, 2, 3, 17, 151]
153 -> [4, 17, 152]
-> [4, 17, 152]
154 -> []
-> [1, 3, 4, 5, 16, 153]


The reason is, that if n - count is not in the result list, you still subtract it when possible.

This is typically a situation for recursion, anticipating the result of a smaller problem:

Speed improvement can be done by considering the square root of tot as next candidate if that is smaller than n - count - 1.