Задача «Теорема Лагранжа» Решение

Теорема Лагранжа утверждает, что любое натуральное число можно представить в виде суммы четырех точных квадратов. По данному числу n найдите такое представление: напечатайте от 1 до 4 натуральных чисел, квадраты которых дают в сумме данное число.

Формат ввода

Программа получает на вход одно натуральное число n < 10000.

Формат вывода

Программа должна вывести от 1 до 4 натуральных чисел, квадраты которых дают в сумме данное число.

Решение

from math import sqrt


def find_squares(n, squares_amount):
    if squares_amount == 1:
        if sqrt(n) == int(sqrt(n)):
            return [int(sqrt(n))]
        return False
    a = 1
    while a * a < n:
        x = find_squares(n - a * a, squares_amount - 1)
        if x:
            return [a] + x
        a += 1
    return False


tries = {0: 1, 1: 1, 2: 2, 3: 3}
number = int(input())
for index in range(tries[number % 4], 5):
    squares = (find_squares(number, index))
    if squares:
        print(*squares)
        break