Представление натурального числа в виде суммы квадратов натуральных чисел, слишком «‎медленный» код

Question

# Представление натурального числа
# в виде суммы квадратов
# натуральных чисел. 
 
k=int(input())
for i in range(1, k+1):
    if k==i**2:
        print(i)
for i in range(1, k+1):
    for j in range(i, k+1):
        if k==i**2+j**2:
            print(i, j)
for i in range(1, k+1):
    for j in range(i, k+1):
        for l in range(j, k+1):
            if k==i**2+j**2+l**2:
                print(i, j, l) 
for i in range(1, k+1):
    for j in range(i, k+1):
        for l in range(j, k+1):
            for m in range(l, k+1):
                if k==i**2+j**2+l**2+m**2:
                    print(i, j, l, m)

Разумеется, такой нерациональный код выполняется крайне медленно, например, на моём домашнем компьютере при вводе числа 2021 программа работала около минуты, тогда как Wolfram Alpha делает то же самое за секунду.

Пожалуйста, подскажите, как изменить код, дабы радикально ускорить данное вычисление.

Answer 1

ну в первом приближении можно в лоб сделать вот такой код:

def alg(value, data):
    if len(data) > 1:
        return
    for i in range(1, int(value**0.5) + 1):
        res = value - i * i
        if res > 0:
            alg(res, data + [i])
        elif res == 0:
            print(' + '.join(f'{j} * {j}' for j in data + [i]), '=', sum(j * j for j in data + [i]))

alg(int(input('введите число')), [])

алгоритм раскладывает на сумму до 2 квадратов, если надо больше, то надо повысить значение вот в этом числе:

if len(data) > 1:

если вообще убрать условие, то будут искаться все возможные суммы, но тогда и 5 = 1 + 1 + 1 + 1 + 1

Недостаток алгоритма - он ищет все комбинации и 5 = 1 + 4 и 5 = 4 + 1. По идее надо от таких вещей избавляться (например, что слагаемые только от меньшего к большему)

Вот более оптимальный алгоритм, который не выдает дубликаты

def alg(value, last = 1, data = []):
    if len(data) > 2:
        return
    for i in range(last, int(value**0.5) + 1):
        res = value - i * i
        if res > 0:
            alg(res, i, data + [i])
        elif res == 0:
            print(' + '.join(f'{j} * {j}' for j in data + [i]), '=', sum(j * j for j in data + [i]))

alg(1256)

Недостаток алгоритма - если вообще убрать проверку на кол-во слагаемых, то питон начинает ругаться на переполнение стека (большая глубина вложенных функций)

Но если в качестве условия потребовать, чтобы слагаемые были уникальными, то тогда это условие можно убрать и код будет таким:

def alg(value, last = 1, data = []):
    for i in range(last + 1, int(value**0.5) + 1):
        res = value - i * i
        if res > 0:
            alg(res, i, data + [i])
        elif res == 0:
            print(' + '.join(f'{j} * {j}' for j in data + [i]), '=', sum(j * j for j in data + [i]))

alg(1256)

Answer 2

Без рекурсий, от проход большего к меньшему. Алгоритм старается использовать в сумме как можно большие числа.

def calc(N):
    sum = 0
    seq = []
    for i in range(int(N**0.5), 0, -1):
            tmp_sum = sum + i*i
            while tmp_sum <= N:
                seq.append(i)
                if tmp_sum == N: 
                    return seq
                sum = tmp_sum
                tmp_sum = sum + i*i

for i in range(1,10000):
    print(i, "=", " + ".join(map(lambda x: f"{x}^2", calc(i))))

Расчёт для всех чисел от нуля до миллиона на моём компе 10-тилетней давности занял 1 минуту 35 секунд.

Answer 3

По теореме Лагранжа о сумме четырёх квадратов четырех слагаемых достаточно для представления любого числа.

Ответ ищется от простого к сложному. Сперва проверяем не является ли само число квадратом (за константу). Затем строим множество квадратов и проверяем можно ли представить число в виде пары слагаемых (за sqrt(N)). Затем строим множество попарных сумм квадратов и ищем представление в виде тройки (за N). Если и это не удалось ищем представление в виде четвёрки (за N).

Общее время работы пропорционально N. Можно сделать намного лучше, но там сложная математика.

import math


def represent(n):
    sq = math.isqrt(n)

    if sq * sq == n:
        return sq,

    set1 = set(i * i for i in range(1, sq + 1))
    for i in set1:
        if n - i in set1:
            return math.isqrt(i), math.isqrt(n - i)

    def gen2():
        for i in set1:
            for j in set1:
                s = i + j
                if s <= n:
                    yield s

    set2 = set(gen2())

    for i in set1:
        if n - i in set2:
            return math.isqrt(i), *represent(n - i)

    for i in set2:
        if n - i in set2:
            return *represent(i), *represent(n - i)

    assert False


n = int(input())
print(n, '=', ' + '.join(f'{i}^2' for i in represent(n)))

$ echo 7 | python represent-n-as-sum-of-squares.py
7 = 1^2 + 1^2 + 1^2 + 2^2

$ echo 2021 | python represent-n-as-sum-of-squares.py
2021 = 16^2 + 1^2 + 42^2

$ time echo 4300007 | python represent-n-as-sum-of-squares.py 
4300007 = 1385^2 + 423^2 + 182^2 + 1473^2

real  0m0.993s
user  0m0.960s
sys   0m0.032s

P.S. В коде есть рекурсия, но она не настоящая. Глубина рекурсии не более двух и все рекурсивные вызовы делаются для чисел, который представимы парой слагаемых, то есть считаются быстро. Можно было бы сделать без рекурсии, но так код короче.

БЛОГ НА HUSL

Представление натурального числа в виде суммы квадратов натуральных чисел, слишком «‎медленный» код

Ответы (3 шт):