Как улучшить скорость функции?

Question

[Yura Khlyan]

Доброе время суток.
У меня задание состоит в следующем: выбрать все пары чисел (a, b) меньших за n, которые удовлетворяют условию: a * b == (sum(1..n) - a - b). Пары (a, b) и (b, a) - разные.
Я написал код, но он слишком медленный. Что можете посоветовать?

def removNb(n):
    return [(a, b)  for a in range(n+1) for b in range(n+1) if sum(range(n+1)) - a - b == a * b]

Comments

[Lynn «Кофеман»]

Ну как минимум не стоит вычислять каждый разс сумму чисел от 1 до n, нет смысла гонять b от 1 до n (можно до a).

А лучше включить мозг и вспомнить математику.
Ваше условие переписывается как (a+1)*(b+1) = (N+1)*N/2 + 1 и задача сводится к поиску разложения на множители числа (N+1)*N/2 + 1

Answer 1

[Mrrl]

Записать условие в виде (a+1)*(b+1) == n*(n+1)/2+1 = N, и для чисел a от (n+1)/2 до n проверить, делится ли N на a+1. Если делится, то b=N/(a+1)-1. На всякий случай, проверить, что b < n+1.
При желании, перебор можно сократить примерно в 12 раз. Но тогда вместо проверки делимости надо будет быстро проверять, является ли число полным квадратом - что может быть не очень просто.

Comments

[bobrovskyserg]

Для проверки на полный квадрат есть всякие хитрости.
Например, глянуть на последние биты числа - квадраты кончаются не чем-попало.
Для последнего байта только 44 значения подходят (как-то писал такую проверку)

[Mrrl]

bobrovskyserg: Я думал обойтись одним шагом метода Ньютона, но его не хватит - там в начале или в конце перебора число, из которого извлекается корень, может очень быстро измениться. Хотя можно проверить для двух самых маленьких чисел с помощью sqrt, а дальше по непрерывности.
В общем, если p*q==N, то (p+q)^2-(p-q)^2==4*N, а p+q у нас лежит между n*sqrt(2) и n*1.5. Перебираем их, проверяем, извлекается ли корень из (p+q)^2-4*N, и если да - печатаем ответ.

Answer 2

[throughtheether]

Что можете посоветовать?

sum(range(n+1)) - константа (в дальнейшем, C) для заданного n, можно вынести ее вычисление за цикл.
Далее, ваше уравнение можно переписать: b = (C-a)/(a+1). Можно перейти от квадратичного времени к линейному, итерируя по a и проверяя делимость (C-a) на (а+1)

Навскидку:

def get_ab(n):
	C=sum(xrange(1,n+1))
	return [(a, (C-a)/(a+1)) for a in xrange(1,n+1) if (C-a)%(a+1)==0 and (C-a)/(a+1)<n]

In [10]: % timeit removNb(10)
1000 loops, best of 3: 185 us per loop
In [11]: %timeit get_ab(10)
100000 loops, best of 3: 6.74 us per loop

Comments

[Yura Khlyan]

После комментария к самому вопросу идентично сам сделал. Хватило скорости.

[throughtheether]

Вы молодец, успехов в учебе.

[Yura Khlyan]

throughtheether:
Спасибо, и Вам хорошего дня.