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

Задача:

задача

Мой код:

def is_simple(value):
    count = 2
    if (value == 2) or ((value != 1) and (value % 2 != 0)):
        for i in range(3, value // 2):
            if value % i == 0:
                count += 1
            if count != 2:
                return False
    return True

def is_sum_simple(value):
    value = str(value)
    sum = 0
    count = 2
    if value != 1:
        for i in range(len(value)):
            sum += int(value[i])
        if (sum == 2) or ((sum != 1) and (sum % 2 != 0)):
            for i in range(3, sum // 2):
                if sum % i == 0:
                    count += 1
                if count != 2:
                    return False
    return True


n = int(input())
t = 0
c = 0
tc = 0
for i in range(2, n + 1):
    if is_simple(i):
        t += 1
    if is_sum_simple(i):
        c += 1
    if is_simple(i) and is_sum_simple(i):
        tc += 1
print(t, c, tc, sep="\n")

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

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

Один из способов упростить код и повысить его эффективность - использовать решето Эратосфена для определения простых чисел.

Также, вместо выполнения проверки is_simple и is_sum_simple для одного числа дважды, можно объединить эти проверки в одну функцию.

Ниже приведен упрощенный код:

def is_prime(n):
    if n < 2:
        return False
    for i in range(2, int(n ** 0.5) + 1):
        if n % i == 0:
            return False
    return True

def is_sum_prime(n):
    total = sum(int(i) for i in str(n))
    return is_prime(total)

n = int(input())

total_primes = 0
total_sum_primes = 0
total_both_primes = 0

for i in range(2, n + 1):
    if is_prime(i):
        total_primes += 1
    if is_sum_prime(i):
        total_sum_primes += 1
    if is_prime(i) and is_sum_prime(i):
        total_both_primes += 1

print(total_primes, total_sum_primes, total_both_primes, sep="\n")
→ Ссылка
Автор решения: Stanislav Volodarskiy

В этой задаче два интересных момента: надо проверять числа на простоту и надо быстро строить суммы цифр для числа. В обоих случаях обрабатываются все числа в диапазоне [0, n], n ≤ 105. Для проверки простоты напрашивается решето Эратосфена. Суммы цифр чисел тоже можно собрать в список, причём для заполнения одной суммы потребуется только одно сложение.

make_sieve(n) создаёт решето Эратосфена [0, n). Решето - массив байт. Нулевой байт соответствует простому числу с этим индексом. Массив байт выбран из-за лучшей константы и меньшего потребления памяти (это связанные вещи).

make_digit_sums(n) заполняет байтовый массив суммами цифр в диапазоне [0, n) и делает это довольно быстро - гораздо быстрее чем считать суммы цифр с помощью sum(map(int, str(i))) в цикле.

Основной цикл бежит по решету и массиву сумм цифр. Решето используется и для проверки простоты сумм цифр.

import math


def make_sieve(n):
    sieve = bytearray(n)
    if n > 0:
        sieve[0] = 1
    if n > 1:
        sieve[1] = 1

    for j in range(4, n, 2):
        sieve[j] = 1

    for i in range(3, math.isqrt(n - 1) + 1, 2):
        if sieve[i] == 0:
            for j in range(i * i, n, 2 * i):
                sieve[j] = 1
    return sieve


def make_digit_sums(n):
    sums = bytearray(n)

    k = 1
    while True:
        m1 = k
        for _ in range(1, 10):
            m2 = min(m1 + k, n)
            for i, j in zip(range(m1, m2), range(m1 - k, m2 - k)):
                sums[i] = 1 + sums[j]
            if m2 == n:
                return sums
            m1 += k
        k = m1


def main():
    n = int(input())

    sieve = make_sieve(n + 1)
    digit_sums = make_digit_sums(n + 1)

    tea = 0
    chakchak = 0
    tea_and_chakchak = 0
    for p, d in zip(sieve, digit_sums):
        if p == 0:
            tea += 1
            if sieve[d] == 0:
                chakchak += 1
                tea_and_chakchak += 1
        elif sieve[d] == 0:
            chakchak += 1

    print(tea)
    print(chakchak)
    print(tea_and_chakchak)


main()

Три миллиона цифр за одну секунду:

$ time echo 3000000 | python chakchak.py
216816
741451
82149

real  0m0.825s
user  0m0.824s
sys   0m0.000s

P.S. Чак-чак обязательно попробую.

P.P.S. Решето можно немного разогнать с помощью техники специфичной для Питона. Если поменять циклы на срезы, можно обработать четыре миллиона чисел за секунду:

def make_sieve(n):
    sieve = bytearray(n)
    if n > 0:
        sieve[0] = 1
    if n > 1:
        sieve[1] = 1

    sieve[4::2] = b'\x01' * ((n - 3) // 2)

    for i in range(3, math.isqrt(n - 1) + 1, 2):
        if sieve[i] == 0:
            sieve[i * i::2 * i] = b'\x01' * ((n - i * (i - 2) - 1) // (2 * i))
    return sieve
→ Ссылка