Генераторы списков

Генераторы списков - это простой и лаконичный способ создания списков из других итерируемых объектов.

Самый простой генератор состоит из выражения и цикла, которые нужно поместить в квадратные скобки:

>>> [x**2 for x in range(10)]
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Сначала следует выражение x**2, которое содержит переменную x, затем следует цикл for x in в котором объявляется таже самая переменная x, а после размещается итерируемый объект range(10). Выражение вовсе не обязано содержать переменные, например, список заполненный нулями можно получить так:

>>> [0 for i in range(10)]
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

Но подобные генераторы-заполнители всегда имеют более простые равносильные команды, например, тот же самый список может быть получен с помощью команды [0]*10

Так же выражение может содержать произвольное количество переменных, но в этом случае для каждой переменной объявляется отдельный цикл и указывается собственный итерируемый объект:

>>> [x + y for x in range(1, 4) for y in range(5, 10)]
[6, 7, 8, 9, 10, 7, 8, 9, 10, 11, 8, 9, 10, 11, 12]

Данный генератор имитирует поведение вложенных циклов for, т.е. список полученный выше может быть создан вот так:

>>> tmp = []
>>> for x in range(1, 4):
...     for y in range(5, 10):
...         tmp.append(x + y)
...
>>> tmp
[6, 7, 8, 9, 10, 7, 8, 9, 10, 11, 8, 9, 10, 11, 12]

Количество итераций внутренних циклов может зависеть от внешних циклов, например:

>>> [str(x) + str(y) for x in range(3) for y in range(x+3)]
['00', '01', '02', '10', '11', '12', '13', '20', '21', '22', '23', '24']

Что аналогично следующему коду:

>>> res = []
>>> for x in range(3):
...     for y in range(x + 3):
...         res.append(str(x) + str(y))
...
>>> res
['00', '01', '02', '10', '11', '12', '13', '20', '21', '22', '23', '24']

Внутренние циклы могут использоваться для итераций по элементам, полученным на внешних циклах, например, для распаковки внутренних списков:

>>> m = [[1], [2, 3], [4, 5, 6], [7, 8, 9]]
>>>
>>> [x for y in m for x in y]
[1, 2, 3, 4, 5, 6, 7, 8, 9]

Аналогичный результат может быть выдан следующим кодом:

>>> res= []
>>> for y in m:
...     for x in y:
...         res.append(x)
...
>>> res
[1, 2, 3, 4, 5, 6, 7, 8, 9]

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

>>> a = [[x, y] for x in range(0, 3) for y in range(0, 3)]
>>> a
[[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2], [2, 0], [2, 1], [2, 2]]

Для создания списка a мы использовали выражение [x, y] результатом которого является список из двух элементов. Поскольку все подсписки являются одинакового размера, то мы можем продемонстрировать как работают генераторы с распаковкой:

>>> [n + m for n, m in a]
[0, 1, 2, 1, 2, 3, 2, 3, 4]

Получить тот же список можно с помощью цикла for:

>>> tmp = []
>>>
>>> for n, m in a:
...     tmp.append(n + m)
...
>>> tmp
[0, 1, 2, 1, 2, 3, 2, 3, 4]

Генераторы списков, поддерживают и более сложный механизм распаковки с помощью оператора *. Для примера снова создадим подходящий итерируемый объект - список a, а затем используем его для создания другого списка с помощью генератора:

>>> a = [[x, y, z] for x in range(1, 4) for y in range(3) for z in range(3)]
>>> a
[[1, 0, 0], [1, 0, 1], [1, 0, 2], [1, 1, 0], [1, 1, 1], [1, 1, 2], [1, 2, 0], [1, 2, 1], [1, 2, 2],
 [2, 0, 0], [2, 0, 1], [2, 0, 2], [2, 1, 0], [2, 1, 1], [2, 1, 2], [2, 2, 0], [2, 2, 1], [2, 2, 2],
 [3, 0, 0], [3, 0, 1], [3, 0, 2], [3, 1, 0], [3, 1, 1], [3, 1, 2], [3, 2, 0], [3, 2, 1], [3, 2, 2]]
>>>
>>> [[n, sum(m)]  for n, *m in a]
[[1, 0], [1, 1], [1, 2], [1, 1], [1, 2], [1, 3], [1, 2], [1, 3], [1, 4],
 [2, 0], [2, 1], [2, 2], [2, 1], [2, 2], [2, 3], [2, 2], [2, 3], [2, 4],
 [3, 0], [3, 1], [3, 2], [3, 1], [3, 2], [3, 3], [3, 2], [3, 3], [3, 4]]

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

>>> n, *m = [1, 0, 0]
>>>
>>> n
1
>>> m
[0, 0]

А работа самого генератора, аналогична следующему циклу:

>>> tmp = []
>>>
>>> for n, *m in a:
...     tmp.append([n, sum(m)])
...
>>> tmp
[[1, 0], [1, 1], [1, 2], [1, 1], [1, 2], [1, 3], [1, 2], [1, 3], [1, 4],
 [2, 0], [2, 1], [2, 2], [2, 1], [2, 2], [2, 3], [2, 2], [2, 3], [2, 4],
 [3, 0], [3, 1], [3, 2], [3, 1], [3, 2], [3, 3], [3, 2], [3, 3], [3, 4]]

Вложенные генераторы

Генераторы списков могут быть вложены друг в друга в результате чего могут быть получены списки списков:

>>> [[x + y for x in 'ABCD'] for y in '123']
[['A1', 'B1', 'C1', 'D1'],
 ['A2', 'B2', 'C2', 'D2'],
 ['A3', 'B3', 'C3', 'D3']]

Данный генератор эквивалентен следующему коду:

>>> res = []
>>> for y in '123':
...     tmp = []
...     for x in 'ABCD':
...         tmp.append(x + y)
...     res.append(tmp)
...
>>> res
[['A1', 'B1', 'C1', 'D1'],
 ['A2', 'B2', 'C2', 'D2'],
 ['A3', 'B3', 'C3', 'D3']]

Можно использовать несколько уровней вложенности:

>>> [[[x + y + z for x in 'ABC'] for y in '123'] for z in '+-*']
[[['A1+', 'B1+', 'C1+'],
  ['A2+', 'B2+', 'C2+'],
  ['A3+', 'B3+', 'C3+']],
  
 [['A1-', 'B1-', 'C1-'],
  ['A2-', 'B2-', 'C2-'],
  ['A3-', 'B3-', 'C3-']],
  
 [['A1*', 'B1*', 'C1*'],
  ['A2*', 'B2*', 'C2*'],
  ['A3*', 'B3*', 'C3*']]]

Данный генератор эквивалентен следующему коду:

>>> res = []
>>>
>>> for z in '+-*':
...     tmp1 = []
...     for y in '123':
...         tmp2 = []
...         for x in 'ABC':
...             tmp2.append(x + y + z)
...         tmp1.append(tmp2)
...     res.append(tmp1)
...
>>> res
[[['A1+', 'B1+', 'C1+'],
  ['A2+', 'B2+', 'C2+'],
  ['A3+', 'B3+', 'C3+']],
  
 [['A1-', 'B1-', 'C1-'],
  ['A2-', 'B2-', 'C2-'],
  ['A3-', 'B3-', 'C3-']],
  
 [['A1*', 'B1*', 'C1*'],
  ['A2*', 'B2*', 'C2*'],
  ['A3*', 'B3*', 'C3*']]]

Вложенные генераторы могут быть переданы функциям:

[sum([x + y for x in range(10)]) for y in range(5)]
[45, 55, 65, 75, 85]

Вложенные генераторы, могут быть использованы для создания различных матриц:

>>> # создаем квадратную матрицу
>>> n = 4
>>> [[i for i in range(j, j + n)] for j in range(0, n**2-1, n)]
[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]
>>>
>>> # создаем прямоугольную матрицу
>>> l, k = 3, 6
>>> [[i for i in range(j, j + k)] for j in range(0, l*k-1, k)]
[[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11], [12, 13, 14, 15, 16, 17]]

С помощью вложенных генераторов можно выполнять более сложные операции над матрицами: транспонирование, отражения, повороты, сложение и т.д. Вот несколько простых примеров:

>>> m = [[i for i in range(j, j + k)] for j in range(0, l*k-1, k)]
>>> m    # исходная матрица
[[0, 1, 2, 3, 4, 5],
 [6, 7, 8, 9, 10, 11],
 [12, 13, 14, 15, 16, 17]]
>>>
>>> # транспонирование m:
>>> [[m[i][j] for i in range(len(m))] for j in range(len(m[0]))]
[[0, 6, 12],
 [1, 7, 13],
 [2, 8, 14],
 [3, 9, 15],
 [4, 10, 16],
 [5, 11, 17]]
>>>
>>> # Поворот m на 90 градусов по часовой стрелке:
>>> [[m[i][j] for i in range(len(m)-1, -1, -1)] for j in range(len(m[0]))]
[[12, 6, 0],
 [13, 7, 1],
 [14, 8, 2],
 [15, 9, 3],
 [16, 10, 4],
 [17, 11, 5]]

Генераторы с условием if

Выполнение действий над переменной может выполняться с произвольным условием if которое указывается после цикла for и проверяется на каждой итерации:

>>> # числа кратные 17:
>>> [x for x in range(100) if x % 17 == 0]
[0, 17, 34, 51, 68, 85]

Работа этого генератора аналогична работе следующего цикла с условием:

>>> res = []
>>> for x in range(100):
...     if x % 17 == 0:
...         res.append(x)
...
>>> res
[0, 17, 34, 51, 68, 85]

Не важно насколько сложным является условие:

>>> [x for x in range(100)
...  if (x**2 % 4 == 1 and x**3 % 8 == 1) or (x**2 % 8 == 1 and x**3 % 4 == 1)]
[1, 5, 9, 13, 17, 21, 25, 29, 33, 37, 41, 45, 49, 53, 57, 61, 65, 69, 73, 77, 81, 85, 89, 93, 97]

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

>>> [int(x + y)
... for x in '0123456789' if x not in '0346789'
... for y in '0123456789' if y != x and y not in '369']
[10, 12, 14, 15, 17, 18, 20, 21, 24, 25, 27, 28, 50, 51, 52, 54, 57, 58]

Получить такой же результат можно с помощью следующего кода:

>>> res = []
>>>
>>> for x in '0123456789':
...     if x not in '0346789':
...         for y in '0123456789':
...             if y != x and y not in '369':
...                 res.append(int(x + y))
...
>>> res
[10, 12, 14, 15, 17, 18, 20, 21, 24, 25, 27, 28, 50, 51, 52, 54, 57, 58]

Если переменных несколько, то все они могут учавствовать в условном выражении, но только после последнего объявления цикла for:

>>> [str(x) + str(y) for x in range(5) for y in range(5) if x*y != 0 and x != y]
['12', '13', '14', '21', '23', '24', '31', '32', '34', '41', '42', '43']

Что аналогично следующему коду:

>>> res = []
>>> for x in range(5):
...     for y in range(5):
...         if x*y != 0 and x != y:
...             res.append(str(x) + str(y))
...
>>> res
['12', '13', '14', '21', '23', '24', '31', '32', '34', '41', '42', '43']

Условные выражения могут использоваться во вложенных генераторах, но необходимо учитывать уровни вложенности, т.е. переменная не может учавствовать в условном выражении, до тех пор пока она не будет извлечена из итерируемого объекта:

>>>  # Это работает:
>>> [[str(x) + str(y) for x in range(4) if x != 2] for y in range(4)]
[['00', '10', '30'], ['01', '11', '31'], ['02', '12', '32'], ['03', '13', '33']]
>>>
>>> # Это не работает:
>>> [[str(x) + str(y) for x in range(4)] for y in range(4) if x != 2]
Traceback (most recent call last):
NameError: name 'x' is not defined

Если нужно использовать условное выражение в котором учавствуют все переменные, то оно должно располагаться в самом внутреннем генераторе:

>>> [[str(x) +str(y) for x in range(10) if x - y > 0] for y in range(10)]
[['10', '20', '30', '40', '50', '60', '70', '80', '90'],
 ['21', '31', '41', '51', '61', '71', '81', '91'],
 ['32', '42', '52', '62', '72', '82', '92'],
 ['43', '53', '63', '73', '83', '93'],
 ['54', '64', '74', '84', '94'],
 ['65', '75', '85', '95'],
 ['76', '86', '96'],
 ['87', '97'],
 ['98'],
 []]

Каждый уровень вложенности может иметь свое условное выражение:

>>> [[[str(x) + str(y) + str(z) for x in range(5) if x + y - z > 0]
...                             for y in range(5) if y - z > 0]
...                             for z in range(5) if z != 0]
[[['021', '121', '221', '321', '421'],
  ['031', '131', '231', '331', '431'],
  ['041', '141', '241', '341', '441']],
             
 [['032', '132', '232', '332', '432'],
  ['042', '142', '242', '342', '442']],
             
 [['043', '143', '243', '343', '443']],
  []]

Генераторы с условием if... else ...

Конструкция if... else ... также может применяться в генераторах, но она должна указываться до цикла for:

>>> [x*4 if x in 'ACEG' else x*2 for x in 'ABCDEFG']
['AAAA', 'BB', 'CCCC', 'DD', 'EEEE', 'FF', 'GGGG']

Данный генератор эквивалентен следующему коду:

>>> res = []
>>> for x in 'ABCDEFG':
...     if x in 'ACEG':
...         res.append(x*4)
...     else:
...         res.append(x*2)
...
>>> res
['AAAA', 'BB', 'CCCC', 'DD', 'EEEE', 'FF', 'GGGG']

Конструкции if... else ... и if могут быть использованы вместе в одном генераторе:

>>> [x**2 if x % 2 == 0 else x/10 for x in range(10) if x % 3 != 0]
[0.1, 4, 16, 0.5, 0.7, 64]

Если в выражении присутствуют несколько переменных, то все они могут быть использованы в конструкции if... else ...:

>>> [str(x) + str(y) if x - y > 0 else str(x)*2 + str(y)*2 for x in range(4) for y in range(4)]
['0000', '0011', '0022', '0033', '10', '1111', '1122', '1133', '20', '21', '2222', '2233', '30', '31', '32', '3333']

даже в случае использования нескольких переменных в генераторе возможно совместное использование конструкций if... else ... и if:

>>> [str(x) + str(y) if x - y > 0 else str(x)*2 + str(y)*2 for x in range(7) if x % 3 != 0 for y in range(4)]
['10', '1111', '1122', '1133', '20', '21', '2222', '2233', '40', '41', '42', '43', '50', '51', '52', '53']

Некоторый интерес представляют случаи совместного использования конструкции if... else ... и вложенных генераторов:

>>> [[x + 10 for x in range(y + 1)] if y % 2 == 0 else [x - 10 for x in range(y + 1)] for y in range(5)]
[[10], [-10, -9], [10, 11, 12], [-10, -9, -8, -7], [10, 11, 12, 13, 14]]

Работа этого генератора может быть выполнена следующим кодом:

>>> for y in range(5):
...     tmp = []
...     if y % 2 == 0:
...         for x in range(y + 1):
...             tmp.append(x + 10)
...     else:
...         for x in range(y + 1):
...             tmp.append(x - 10)
...     res.append(tmp)
...
>>> res
[[10], [-10, -9], [10, 11, 12], [-10, -9, -8, -7], [10, 11, 12, 13, 14]]

Используя условия внутри генераторов важно избегать повторяющихся операций, например:

>>> >>> [x + y for x in range(5) for y in range(5) if (x + y) % 2 == 0]
[0, 2, 4, 2, 4, 2, 4, 6, 4, 6, 4, 6, 8]

В Python версии \(3.8\) появился оператор := который позволяет выполнять одинаковые операции только в одном месте, благодаря чему приведенный выше генератор может быть записан как:

>>> [z for x in range(5) for y in range(5) if (z := x + y) % 2 == 0]
[0, 2, 4, 2, 4, 2, 4, 6, 4, 6, 4, 6, 8]

Таже рекомендация касается повторного использования ресурсоемких функций, например:

>>> def rand(n):
    r = 43
    for i in range(100000 + n):
        r = int(str(int(r**0.5)**r + r)[:2])
    return r
>>>
>>> [rand(x) for x in range(20) if rand(x) % 2 != 0]
[25, 29, 25, 29, 25, 29, 25, 29, 25, 29]

В данном случае так же целесообразнее воспользоваться оператором :=

>>> [z for x in range(20) if (z := rand(x)) % 2 != 0]
[25, 29, 25, 29, 25, 29, 25, 29, 25, 29]

Если вы используете Python версии ниже \(3.8\), то что бы избежать двойного выполнения операций можете обеспечить их выполнение во внутреннем генераторе, а во внешнем использовать его результат:

>>> [z for z in [x + y for x in range(5) for y in range(5)] if z % 2 == 0]
[0, 2, 4, 2, 4, 2, 4, 6, 4, 6, 4, 6, 8]
>>>
>>> [z for z in [rand(x) for x in range(20)] if z % 2 != 0]
[25, 29, 25, 29, 25, 29, 25, 29, 25, 29]

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

>>> [z for x in range(5) for y in range(5) for z in [x + y] if z % 2 == 0]
[0, 2, 4, 2, 4, 2, 4, 6, 4, 6, 4, 6, 8]
>>>
>>> [z for x in range(20) for z in [rand(x)] if z % 2 != 0]
[25, 29, 25, 29, 25, 29, 25, 29, 25, 29]