Найти число, сумма квадратов делителей которого является квадратом целого числа

Нужно найти все такие числа в интервале от 1 до 250.

РЕШЕНИЕ РАБОТАЕТ, просто нужно оптимизировать

codewars не принимает из за времени выполнения(

import time

t1 = time.time()
def list_squared(m, n):
    need_numbers = []
    for i in range(m, n):
        divisors = [elem*elem for elem in range(1, (i+1)/elem) if i % elem == 0 ]       
        da = sum(divisors)**(0.5)
        if da.is_integer():
            need_numbers.append([i, sum(divisors)]) 
    return need_numbers



print(list_squared(1, 250))
t2 = time.time()
print(t2-t1)

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

Автор решения: Pak Uula

Уберите деление на elem в списке, оставьте range(1, (i+1)), и тогда всё заработает.

Время счёта 1.5 миллисекунды. Мне кажется, оптимизировать нет необходимости.

import time

t1 = time.time()
def list_squared(m, n):
    need_numbers = []
    for i in range(m, n):
        divisors = [elem*elem for elem in range(1, (i+1)) if i % elem == 0 ]       
        da = sum(divisors)**(0.5)
        if da.is_integer():
            need_numbers.append([i, sum(divisors)]) 
    return need_numbers



print(list_squared(1, 250))
t2 = time.time()
print(t2-t1)

Вывод:

[[1, 1], [42, 2500], [246, 84100]]
0.0015308856964111328

UPDATE

Топик-стартер попросил оптимизированное решение. Ок.

Оптимизированное решение использует тот факт, что все делители числа n состоят из тех же простых чисел, что делят само число n. Например, делители числа 300 1, 5, 25, 3, 15, 75, 2, 10, 50, 6, 30, 150, 4, 20, 100, 12, 60, 300 сами делятся только на 2,3,5. Ну и единица, понятное дело.

Поэтому для ускорения счёта нужно разложить число n на простые множители, найти для каждого множителя с какой степенью оно делит число n и построить все комбинации произведений простых сомножителей.

Пример 300 = 2**2 * 3 * 5**2 Значит, делителями будут: 1, 2**1, 2**2, 3, 2*3, 2**2 * 3, 5, 2**2 * 5, 2*3*5, 2**2 *3*5, 5**2, 2**2 * 5**2, 2*3*5**2, 2**2 *3*5**2, 3*5, 3 * 5**2.

import math

# Список простых чисел, меньших 1000
primes = [2, 3, 5, 7, 11, 13, 17, 19, 23,
29, 31, 37, 41, 43, 47, 53, 59, 61, 67,
71, 73, 79, 83, 89, 97, 101, 103, 107, 109,
113, 127, 131, 137, 139, 149, 151, 157, 163, 167,
173, 179, 181, 191, 193, 197, 199, 211, 223, 227,
229, 233, 239, 241, 251, 257, 263, 269, 271, 277,
281, 283, 293, 307, 311, 313, 317, 331, 337, 347,
349, 353, 359, 367, 373, 379, 383, 389, 397, 401,
409, 419, 421, 431, 433, 439, 443, 449, 457, 461,
463, 467, 479, 487, 491, 499, 503, 509, 521, 523,
541, 547, 557, 563, 569, 571, 577, 587, 593, 599,
601, 607, 613, 617, 619, 631, 641, 643, 647, 653,
659, 661, 673, 677, 683, 691, 701, 709, 719, 727,
733, 739, 743, 751, 757, 761, 769, 773, 787, 797,
809, 811, 821, 823, 827, 829, 839, 853, 857, 859,
863, 877, 881, 883, 887, 907, 911, 919, 929, 937,
941, 947, 953, 967, 971, 977, 983, 991, 997]

def find_prime_divisor(n):
    "Функция наименьший простой делитель числа n"
    lim = math.floor(math.sqrt(n+1))
    # Сначала проверим простые из списка
    for p in primes:
        if p > lim:
            # p*p > n -- дальше перебирать бессмысленно, n не делится на такое большое p
            break
        if n%p == 0:
            # Нашли делитель
            return p
    if p > lim:
        # n меньше миллиона, и не делится ни на одно простое из списка
        # следовательно, n - простое
        return n
    # n больше миллиона. Ищем его делители лобовым перебором. 
    for m in range(p, n,2):
        if n%m == 0:
            return m
    # делители не найдены. n - простое число
    return n

def get_divisor_degree(m,n):
    """Функция возвращает максимальную степень d числа m, при которой m**d делит n
    Возвращается пара (d, n/m**d)
    """
    deg = 0
    while n%m == 0:
        deg += 1
        n //= m
    return deg, n

def factorize(n):
    """Функция раскладывает n на простые множители.
    Возвращается набор пар (простой_делитель, степень_делителя) в виде словаря."""
    factors = {}
    while n > 1:
        p = find_prime_divisor(n)
        deg, n = get_divisor_degree(p,n)
        factors[p] = deg
    return factors

def gen_divisors(factors):
    """Генератор делителей из списка пар (простой_делитель, степень_делителя)"""
    if isinstance(factors, dict):
        factors = list(factors.items())
    if len(factors) == 0:
        # Пустой список - вовзращаем 1
        yield 1
    else:
        # вынимает первый делитель из списка
        p, deg = factors[0]
        p_m = 1
        for _ in range(deg+1):
            # генерируем делители из остальных простых делителей 
            # и умножаем их последовательно на степени p
            for div in gen_divisors(factors[1:]):
                yield p_m*div
            p_m *= p
​
def list_divisors(n):
    '''Функция возвращает список всех делителей числа n'''
    return list(gen_divisors(factorize(n)))
​
def is_square(n):
    "Функция возвращает True, если целое число n является квадратом другого целого числа."
    root = math.floor(math.sqrt(n+1))
    return  root*root == n
is_square(625), is_square(101)
(True, False)

def check_number(n):
    "Функция возвращает True если число подходит под условие задачи, и сумму квадратов делителей"
    divs = list_divisors(n)
    divs_square = sum(map(lambda x:x*x, divs))
    return is_square(divs_square), divs_square

def run(iterable):
    "Функция проверяет каждое число из итератора и генерирует пары (число, сумма квадратов делителей) для подходящих" 
    for n in iterable:
        good, divs_square = check_number(n)
        if good:
            yield n, divs_square

print(list(run(range(1, 10000))))

Вывод:

[(1, 1), (42, 2500), (246, 84100), (287, 84100), (728, 722500), (1434, 2856100), (1673, 2856100), (1880, 4884100), (4264, 24304900), (6237, 45024100), (9799, 96079204), (9855, 113635600)]

Время работы 130-150 миллисекунд. Решение в лоб я не дождался.

Jupyter Notebook с решением

Оптимизированное решение проигрывает решению в лоб на интервале [1, 500] и твёрдо выигрывает на интервалах длиннее чем [1,1000]

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

baseline

Я переделал программу из вопроса. Алгоритм остался тем же.

import math

for n in range(*map(int, input().split())):
    s2 = sum(d * d for d in range(1, n + 1) if n % d == 0)
    sqrt = math.isqrt(s2)
    if sqrt * sqrt == s2:
        print(n, s2)

Сложность O(N(N - M)) для диапазона [M, N). Для диапазона вида [1, N), сложность квадратичная: O(N2). Парабола хорошо видна: baseline graph

σ2

Сумма квадратов делителей числа - частный случай функции делителей. Её можно вычислить по разложению числа на простые множители. Сложность не превосходит O(N√N). Первый множитель - количество чисел, второй - время проверки одного числа.

import math


def factors(n):
    i = 2
    di = 1
    n_sqrt = math.isqrt(n)
    while i <= n_sqrt:
        if n % i == 0: # i is prime divisor of n
            n //= i
            e = 1
            while n % i == 0:
                n //= i
                e += 1
            yield i, e
            n_sqrt = math.isqrt(n)
        i += di
        di = 2

    if n > 1: # n is prime
        yield n, 1


def sigma2(n):
    s = 1
    for p, e in factors(n):
        s *= (p ** (2 * e + 2) - 1) // (p * p - 1)
    return s


def main():
    for n in range(*map(int, input().split())):
        s = sigma2(n)
        s_sqrt = math.isqrt(s)
        if s_sqrt * s_sqrt == s:
            print(n, s)


main()

sigma-factors graph

Решето σ2

Считать разложения соседних чисел расточительно. Например, мы проверяем все числа на делимость на семь. Но на семь делится каждое седьмое число диапазона, не нужно делать столько делений. Новая функция factors разлагает сразу много чисел, экономя время. Время разложения O(NlogN). К сожалению, новая функция требует память O(N), старые версии работали в константной памяти:

import math


def divisors(n):
    yield 2
    yield from range(3, n, 2)


def factors(n):
    ns = list(range(n))

    for p in divisors(math.isqrt(n - 1) + 1):
        if ns[p] == p: # p is prime
            for i in range(p, n, p):
                assert ns[i] % p == 0
                ns[i] //= p
                e = 1
                while ns[i] % p == 0:
                    ns[i] //= p
                    e += 1
                yield i, p, e

    for i in range(n):
        if ns[i] > 1:
            yield i, ns[i], 1


def sigma2(n):
    ss = [1] * n
    for i, p, e in factors(n):
        ss[i] *= (p ** (2 * e + 2) - 1) // (p * p - 1)
    return ss


def main():
    for n, s in enumerate(sigma2(int(input()))):
        s_sqrt = math.isqrt(s)
        if s_sqrt * s_sqrt == s and n > 0:
            print(n, s)


main()

Не забудем про константу. Если совместить разложение на простые и вычисление σ2, можно ускорится почти в два раза:

import math


def divisors(n):
    yield 2
    yield from range(3, n, 2)


def sigma2(n):
    ss = [1] * n
    ns = list(range(n))

    for p in divisors(math.isqrt(n - 1) + 1):
        if ns[p] == p: # p is prime
            p2 = p * p
            for i in range(p, n, p):
                ns[i] //= p
                m = p2
                f = 1 + p2
                while ns[i] % p == 0:
                    ns[i] //= p
                    m *= p2
                    f += m
                ss[i] *= f

    for i, n in enumerate(ns):
        if n > 1:
            ss[i] *= 1 + n * n

    return ss


def main():
    for n, s in enumerate(sigma2(int(input()))):
        s_sqrt = math.isqrt(s)
        if s_sqrt * s_sqrt == s and n > 0:
            print(n, s)


main()

sieve-factors and sieve graphs

Сегментированное решето σ2

Линейная память для решета - заметный недостаток. Решето можно порезать на сегменты. Наилучший размер сегмента √N. Кажется, что введение сегментов должно замедлить рассчёт, но этого не происходит: накладные расходы малы, а маленькое решето лучше использует кеш процессора. Сложность алгоритма по прежнему O(NlogN), памяти нужно меньше - O(√N). Ограничение по памяти не строгое: если N очень велико, а памяти мало, можно уменьшить размер сегмента. Программа замедлится, но продолжит считать.

import math


def divisors(n):
    yield 2
    yield from range(3, n, 2)


def factors(n1, n2):
    ns = list(range(n1, n2))

    for p in divisors(math.isqrt(n2 - 1) + 1):
        i1 = -n1 % p # (n1 + i1) % p == 0
        if i1 < n2 - n1 and ns[i1] % p == 0: # ... and p is prime 
            for i in range(i1, n2 - n1, p):
                ns[i] //= p
                e = 1
                while ns[i] % p == 0:
                    ns[i] //= p
                    e += 1
                yield n1 + i, p, e

    for i, n in enumerate(ns, start=n1):
        if n > 1:
            yield i, n, 1


def sigma2(n1, n2):
    ss = [1] * (n2 - n1)
    for i, p, e in factors(n1, n2):
        ss[i - n1] *= (p ** (2 * (e + 1)) - 1) // (p * p - 1)
    return enumerate(ss, start=n1)


def main():
    n1, n3 = map(int, input().split())
    while n1 < n3:
        n2 = min(n1 + math.isqrt(n1), n3)
        for i, s in sigma2(n1, n2):
            s_sqrt = math.isqrt(s)
            if s_sqrt * s_sqrt == s:
                print(i, s)
        n1 = n2


main()

И снова можно ускорится в два раза, если объединить факторизацию и вычисление функции делителей:

import math


def divisors(n):
    yield 2
    yield from range(3, n, 2)


def sigma2(n1, n2):
    ns = list(range(n1, n2))
    ss = [1] * (n2 - n1)

    for p in divisors(math.isqrt(n2 - 1) + 1):
        i1 = -n1 % p # (n1 + i1) % p == 0
        if i1 < n2 - n1 and ns[i1] % p == 0: # ... and p is prime 
            p2 = p * p
            for i in range(i1, n2 - n1, p):
                ns[i] //= p
                t = p2
                f = 1 + p2
                while ns[i] % p == 0:
                    ns[i] //= p
                    t *= p2
                    f += t
                ss[i] *= f

    for i, n in enumerate(ns):
        if n > 1:
            ss[i] *= 1 + n * n

    return enumerate(ss, start=n1)


def main():
    n1, n3 = map(int, input().split())
    while n1 < n3:
        n2 = min(n1 + math.isqrt(n1), n3)
        for i, s in sigma2(n1, n2):
            s_sqrt = math.isqrt(s)
            if s_sqrt * s_sqrt == s:
                print(i, s)
        n1 = n2


main()

segmented-sieve-factors and segmented-sieve graphs

Всё вместе

Все времена работы на одном логарифмическом графике. Кроме самих времен добавлена экстраполяция до одного часа. За одну минуту самый быстрый вариант обрабатывает примерно в тысячу раз больше чисел чем самый медленный. Для одного часа разница примерно в десять тысяч раз.

all graphs log-log

→ Ссылка