March 16, 2024

List comprehension и не только

List comprehension можно перевести по разному, будь то генератор списков, списочное выражение или списочное включение, но я буду использовать предпочтительно оригинал "list comprehension" и "списочное выражение". :)

List comprehension или списочные выражения – это простой для чтения, компактный и удобный способ создания и обработки списка из любого существующего итерируемого объекта.

Обычно это всего одна строка кода, заключенная в квадратных скобках. List comprehension можно использовать для фильтрации, форматирования, изменения или выполнения других небольших задач с существующими итерируемыми объектами, такими как строки, кортежи, множества, списки и т.д.

Разберем несколько способов создания списочных выражений, а так же выражений со словарями и множествами.

Синтаксис и первый list comprehension

По сути, синтаксис состоит из трех компонентов:

  • результат (expression_out)
  • цикл for (for item in iterable)
  • условие (if condition)
[expression_out for item in iterable if condition]

Пример простого списочного выражения в сравнении с циклом for

new_list = [element for element in range(1, 10)]
print(new_list)
> [1, 2, 3, 4, 5, 6, 7, 8, 9]

# OR

new_list = []
for element in range(1, 10):
    new_list.append(element)
print(new_list)
> [1, 2, 3, 4, 5, 6, 7, 8, 9]

Мы создали новый список всего в одну строку, тогда как на этот простой пример с помощью цикла for ушло бы 3.

Также мы можем использовать любое выражение для изменения элементов списка, например:

new_list = [element + 1 for element in range(1, 10)]
print(new_list)
> [2, 3, 4, 5, 6, 7, 8, 9, 10]

new_list = [element - 1 for element in range(1, 10)]
print(new_list)
> [0, 1, 2, 3, 4, 5, 6, 7, 8]

new_list = [element ** 2 for element in range(1, 10)]
print(new_list)
> [1, 4, 9, 16, 25, 36, 49, 64, 81]

List comprehension с одиночным условием if

В списочное выражение также можно добавить if-условие, которое может помочь отфильтровать данные.

Например, в приведенном ниже коде мы сохраняем в новый список все значения из последовательности, возведенные в квадрат, если они меньше 5:

new_list = [element ** 2 for element in range(1, 10) if element < 5]
print(new_list)
> [1, 4, 9, 16]

# OR

new_list = []
for element in range(1, 10):
    if element < 5:
        new_list.append(element ** 2)

print(new_list)
> [1, 4, 9, 16]

List comprehension с условиями if-else

Предположим, нам надо преобразовать последовательность от 1 до 10 так, чтобы возвести каждое число в степень при условии, что это число меньше 5, в противном случае записываем вместо этого числа 0:

new_list = [element ** 2 if element < 5 else 0 for element in range(1, 10)]
print(new_list)
> [1, 4, 9, 16, 0, 0, 0, 0, 0]

# OR

new_list = []
for element in range(1, 10):
    if element < 5:
        new_list.append(element ** 2)
    else:
        new_list.append(0)

print(new_list)
> [1, 4, 9, 16, 0, 0, 0, 0, 0]

List comprehension с условиями и вложенными циклами

Когда мы работаем с вложенной структурой данных, ее надо и обрабатывать соответствующим образом. Рассмотрим пример на списке из списков и попробуем:

  • выпрямить этот список
  • преобразовать каждое значение из int в str, если оно > 5 и добавить текст "_str"
  • иначе, если условие не соблюдается, умножаем число на 2 и оставляем его типа int

Задание не нагружено каким-либо смыслом, но для примера будет интересно:

base_list = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

new_list = [f"{element}_str" if element > 5 else element * 2 for sublist in base_list for element in sublist]
print(new_list)
> [2, 4, 6, 8, 10, '6_str', '7_str', '8_str', '9_str']

# OR

new_list = []
for sublist in base_list:
    for element in sublist:
        if element > 5:
            new_list.append(f"{element}_str")
        else:
            new_list.append(element * 2)
print(new_list)
> [2, 4, 6, 8, 10, '6_str', '7_str', '8_str', '9_str']

Наш код обрабатывает уже условия поинтереснее, а написан также в одну строку.

Какие еще бывают выражения в Python?

Подобный принцип мы можем распространить на словари и множества. Рассмотрим по 1 примеру:

Dict comprehension

new_dict = {str(key): value * 2 for key, value in enumerate(range(5))}
print(new_dict)
> {'0': 0, '1': 2, '2': 4, '3': 6, '4': 8}

Что происходит: мы создаем словарь из последовательности чисел от 0 до 4 (включительно), где ключ преобразуем в строку, значение умножаем на 2. enumerate используем для получения пары индекс-значение элемента, индекс используем как ключ словаря.

Set comprehension

new_set = {value * 2 for value in [1, 1, 2, 3, 3, 3]}
print(new_set)
> {2, 4, 6}

Set comprehension позволяет проделывать то же самое, что и list comprehension, но только для уникальных элементов. В примере мы избавились от дубликатов и умножили уникальные числа на 2.

Выводы

Списочные выражения не только более компактны, но также очень эффективны. Они будут работать быстрее цикла for.

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