Оптимизация. Количество возможных решений Python

Имеется задание, найти количество решений, нужно оптимизировать , чтобы этот код выполнялся используя не 4 цикла , а 3 (но лучше 2 цикла). Нужно понизить сложность кода.

Формула:

xa^2 + yb^2 = xc^2 + yd^2
n , x , y = map(int,input().split())

counter = 0
for a in range(n+1):
    for b in range(n+1):
        for c in range(n+1):
            for d in range(n+1):
                if (x*(a**2))+(y*(b**2)) == (x*(c**2)) +(y*(d**2)):
                    counter += 1
print(counter)

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

Автор решения: Zhihar

вариант 1 (быстрее исходного):

counter = 0

arr = []
for a in range(n + 1):
    for b in range(n + 1):
        arr.append(a**2 - b**2)

for value1 in arr:
    for value2 in arr:
        if x * value1 == y * value2:
            counter += 1

print(counter)

вариант 2 (еще быстрее):

counter = 0

arr = []
for a in range(n + 1):
    for b in range(n + 1):
        arr.append(a**2 - b**2)

for value1 in arr:
    value2 = x * value1 // y
    counter += arr.count(value2)

print(counter)

вариант 3 (возможно чуть медленнее, чем вариант 2):

arr = [(a**2 - b**2) for a in range(n + 1) for b in range(n + 1)]

counter = sum(arr.count(x * value // y) for value in arr)

print(counter)
→ Ссылка
Автор решения: Danis
list_ = [x * a * a + y * b * b
          for a in range(n + 1)
          for b in range(n + 1)]


counter = 0
for a in set(list_):
    counter += list_.count(a) ** 2
print(counter)

Правую и левую часть можно вычислять заранее

left  = [x * i * i for i in range(n + 1)]
right = [y * i * i for i in range(n + 1)]
list_ = [a + b
         for a in left
         for b in right]

counter = 0
for a in set(list_):
    counter += list_.count(a) ** 2
print(counter)

list.count очень не эффективен, быстрее будет пройтись по всему списку и добавлять значения в словарь

left  = [x * i * i for i in range(n + 1)]
right = [y * i * i for i in range(n + 1)]
list_ = [a + b
         for a in left
         for b in right]

dict_ = {el: 0 for el in set(list_)}

for el in dict_:
    dict_[el] += 1
   
print(sum(el * el for el in dict_.values()))
→ Ссылка
Автор решения: MarianD

Оптимизированный вариант вашей программы — около 5 раз быстрее:

n, x, y = map(int, input().split())

counter = 0
for a in range(n+1):
    a = x * a * a
    for b in range(n+1):
        b = y * b * b
        left = a + b
        for c in range(n+1):
            c = x * c * c
            for d in range(n+1):
                d = y * d * d
                if left == c + d:
                    counter += 1
print(counter)

Объяснение:

Выражения x*a^2, y*b^2 (и их сумма), x*c^2, y*d^2 вычисляются только один раз.

→ Ссылка
Автор решения: Zhihar

У @Danis получилось очень хорошее и очень быстрое решение, но его можно сделать еще быстрее (ориентировочно где-то в 5-10 раз - проверял на n = 200, x = 10, y = 20):

list_ = [x * a * a + y * b * b
          for a in range(n + 1)
          for b in range(n + 1)]

list_.sort()
list_.append(list_[-1] + 1)

counter = 0
local = 1
for i in range(1, len(list_)):
    if list_[i - 1] == list_[i]:
        local += 1
    else:
        counter += local**2
        local = 1

print(counter)

Вот еще более быстрый (где-то на 5%) вариант:

list_ = [x * a * a + y * b * b
          for a in range(n + 1)
          for b in range(n + 1)]

list_.sort()
list_.append(list_[-1] + 1)

counter = 0
pos = 0
for i in range(1, len(list_)):
    if list_[i - 1] != list_[i]:
        counter += (i - pos)**2
        pos = i

print(counter)

Проверял на x = 10, y = 20, n = 2000, время выполнения 1,92сек

Вот этот код работает еще на 5% быстрее:

list_ = [x * a * a + y * b * b
          for a in range(n + 1)
          for b in range(n + 1)]

list_.sort()
list_.append(list_[-1] + 1)

counter = 0
pos = 0
value = list_[0]
for i in range(1, len(list_)):
    if list_[i] != value:
        counter += (i - pos)**2
        pos = i
        value = list_[i]

print(counter)

Вот этот код работает еще на 18% быстрее:

list_ = []

for a in range(n + 1):
    value = x * a * a
    for b in range(n + 1):
        list_.append(value + y * b * b)

list_.sort()
list_.append(list_[-1] + 1)

counter = 0
pos = 0
value = list_[0]

for i in range(1, len(list_)):
    if list_[i] != value:
        counter += (i - pos)**2
        pos = i
        value = list_[i]

print(counter)

Вот этот код работает еще на 13% быстрее:

abx = []
aby = []

for i in range(n + 1):
    abx.append(x * i * i)
    aby.append(y * i * i)

list_ = []

for a in abx:
    for b in aby:
        list_.append(a + b)

list_.sort()
list_.append(list_[-1] + 1)

counter = 0
pos = 0
value = list_[0]

for i in range(1, len(list_)):
    if list_[i] != value:
        counter += (i - pos)**2
        pos = i
        value = list_[i]

print(counter)

Вот этот код работает ещё на 22% быстрее (спасибо @Danis за идею уйти от степеней):

abx = []
aby = []

for i in range(n + 1):
    abx.append(x * i * i)
    aby.append(y * i * i)

list_ = []

for a in abx:
    for b in aby:
        list_.append(a + b)

list_.sort()
list_.append(list_[-1] + 1)

counter = 0
pos = 0
value = list_[0]

for i in range(1, len(list_)):
    if list_[i] != value:
        counter += (i - pos) * (i - pos)
        pos = i
        value = list_[i]

print(counter)

Вот этот код работает еще на 9% быстрее:

abx = []
aby = []

for i in range(n + 1):
    abx.append(x * i * i)
    aby.append(y * i * i)

list_ = [a + b for a in abx for b in aby]

list_.sort()
list_.append(list_[-1] + 1)

counter = 0
pos = 0
value = list_[0]

for i in range(1, len(list_)):
    if list_[i] != value:
        counter += (i - pos) * (i - pos)
        pos = i
        value = list_[i]

print(counter)
→ Ссылка
Автор решения: Exe

Вы также можете использовать библиотеку numba, которая использует just in time компилятор(дает ускорение примерно в 40 раз), в ней можно отключать GIL и использовать мультипоточность, также есть поддержка cuda. Использование очень простое, накидываем декоратор njit на функцию и добавляем параметры какие захотим и prange есть для распаралеливания.

from numba import njit

@njit(nogil=True,parallel=True)   
def foo():
    pass
→ Ссылка
Автор решения: MarianD

Слегка улучшен конец красивого решения Danis-а — применением класса Counter:

from collections import Counter

n, x, y = map(int, input().split())

list_ = [x * a * a + y * b * b
          for a in range(n + 1)
          for b in range(n + 1)]

counter = sum(val * val for val in Counter(list_).values())
print(counter)
→ Ссылка