Подсчёт значений генератора c условием

Иногда требуется подсчитать количество значений генератора, которые отвечают какому-то условию. Есть два способа это сделать: sum(1 for ... if условие) и sum(условие for ...). Второй пользуется тем, что в арифметике False отображается в ноль, а True в единицу. Так что оба варианта считают одно и то же. Какой предпочитаете вы и почему?

Например тривиальный способ подсчёта количества делителей числа:

n = 100

s1 = sum(1 for i in range(1, n + 1) if n % i == 0)

s2 = sum(n % i == 0 for i in range(1, n + 1))

P.S. Вопрос про язык, не про алгоритм. Не предлагайте сосчитать делители по другому, быстрее.


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

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

По моим подсчётам получается второй вариант работает на 20% медленнее (при данном конкретном условии отсева чисел).

import time

n = 10_000_000

t = time.time()
s1 = sum(1 for i in range(1, n + 1) if n % i == 0)
print(time.time() - t)
# 1.008479356765747

t = time.time()
s2 = sum(n % i == 0 for i in range(1, n + 1))
print(time.time() - t)
# 1.2111566066741943

Я даже могу предположить - почему.

  • Число циклов - одинаковое
  • Проверка на остаток от деления - одинаковое число раз

Что же разное? А разное - количество элементов, которые выдаёт генератор, который дёргается функцией sum через итератор. Вот, видимо, большее кол-во элементов, пропускаемое через yield (и через sum!) и даёт некоторое замедление.

Я как обычно и через dis посмотрел.

Кусочек от первого варианта:

              8 LOAD_GLOBAL              0 (n)
             10 LOAD_FAST                1 (i)
             12 BINARY_MODULO
             14 LOAD_CONST               0 (0)
             16 COMPARE_OP               2 (==)
             18 POP_JUMP_IF_FALSE        2 (to 4)
             20 LOAD_CONST               1 (1)
             22 YIELD_VALUE

Кусочек от второго варианта:

              6 STORE_FAST               1 (i)
              8 LOAD_GLOBAL              0 (n)
             10 LOAD_FAST                1 (i)
             12 BINARY_MODULO
             14 LOAD_CONST               0 (0)
             16 COMPARE_OP               2 (==)
             18 YIELD_VALUE

Как видим, в первом варианте действительно YIELD_VALUE работает не на каждый элемент, а во втором на каждый.

P.S. Если почти все элементы в итоге пройдут через фильтр, то первый вариант будет работать медленнее из-за "лишней" инструкции перехода по условию POP_JUMP_IF_FALSE (переход при этом не происходит, но тестируется условие перехода). По идее эта разница не должна быть большой. Но тесты показывают 6-7% разницы.

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

Я буду вести счёт. Начнём:

Краткость (0-1)

Второй вариант короче. Если вы спешите это важно. Это объективный критерий.

Ясность (0.5-1.5)

Оба варианта примерно наравне.

sum(1 for i in range(1, n + 1) if n % i == 0). Зачем складывать единички? Чтобы что-то сосчитать? Обычно для подсчётов есть специальная функция: что-нибудь с названием count. Но в Питоне она отсутствует, приходится считать единицы. Надо привыкнуть.

sum(n % i == 0 for i in range(1, n + 1)) тоже требует напряжения ума: надо понять что результат выражения булев, знать что булевы значения отображаются в целые числа так что сумма этих чисел будет равна количеству истинных значений. В общем ничья.

Хрупкость кода (1.5-1.5)

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

def s1(seq):
    return sum(1 for s in seq if s)

def s2(seq):
    return sum(s for s in seq)


print(s1(['', 'a', 'ab']))
print(s2(['', 'a', 'ab']))

Первый вариант сработает, второй упадёт при попытке сложить число и строку. Его надо исправить
так: sum(bool(s) for s in seq) или
так: sum(s != '' for s in seq) или
так: sum(len(s) != 0 for s in seq) или
так: sum(map(bool, seq)).

Или если вы решите написать функцию count для предиката:

def count1(seq, pred):
    return sum(1 for v in seq if pred(v))


def count2(seq, pred):
    return sum(pred(v) for v in seq)

Второй вариант снова более хрупкий. Питон разрешает в логических выражениях вместо True/False использовать другие значения. Например 42 or x != y вернёт 42. Так что count2 лучшее переписать
так: sum(bool(pred(v)) for v in seq) или
так: sum(map(lambda v: bool(pred(v)), seq)).

Производительность (2-2)

В соседнем ответе производительность разобрана на уровне ассемблера. Поставьте ему плюсик.

NB Я измерял на CPython 3.12. На других версиях результаты могут различаться.

Если коротко, то второй вариант тратит время складывая нули, первый вариант тратит время на условный переход, которого нет во втором. В среднем первый быстрее, но преимущество растворяется на фоне времени вычисления предиката.

Первый вариант быстрее, если количество True мало, а последовательность длинна:

from time import perf_counter

n = 100_000_000

start = perf_counter()
print(sum(1 for i in range(1, n + 1) if n % i == 0), perf_counter() - start)

start = perf_counter()
print(sum(n % i == 0 for i in range(1, n + 1)), perf_counter() - start)

Первый вариант на 28% быстрее второго:

$ python benchmark.py
81 5.596365669043735
81 7.764751962851733

Напротив, первый вариант медленнее если почти все значения – True:

from time import perf_counter

n = 100_000_000

start = perf_counter()
print(sum(1 for i in range(1, n + 1) if n % i != 0), perf_counter() - start)

start = perf_counter()
print(sum(n % i != 0 for i in range(1, n + 1)), perf_counter() - start)

Теперь уже второй вариант быстрее первого на 7%:

$ python benchmark.py
99999919 8.247191620990634
99999919 7.672776157036424

Я разделил очко поровну между вариантами, хотя первый имеет преимущество, пусть не очень большое. Но для кого-то это важно, имейте в виду.

Итоги

По очкам паритет, но если бы у меня было дополнительное очко, я бы отдал его первому варианту. Краткость меня мало волнует, а вот хрупкость второго варианта может приводить к неприятным ошибкам. И производительность в среднем немного лучше у первого варианта.

→ Ссылка