Как выпрямить (flatten) список в Python

автор рубрика ,
Как выпрямить (flatten) список в Python

Списки в Python используются постоянно, а иногда приходится работать с вложенными списками. В этой статье даны некоторые способы выравнивания (flatten) списков, когда вложенные списки выводят на один уровень вверх.

В чем заключается проблема

Возможно вы встречали такого рода вложенные списки (list), т.е. такие списки, которые сами содержат списки. А ваша задача вынести вложенные элементы в один единственный список. Например, исходный список в Python такой:

groups = [['Hong', 'Ryan'], ['Andry', 'Ross'], ['Mike', 'Smith']]

Нужно его преобразовать следующим образом:

['Hong', 'Ryan', 'Andry', 'Ross', 'Mike', 'Smith']

Задача сводится к выравниванию списка на один уровень. С другой стороны, распаковка списка, который содержит список, который в свою очередь также содержит список и т.д. подразумевает распаковку на n-й уровень, где n — уровень вложенности.

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

Выравнивание с помощью циклов for

Самый простой способ — это использование цикла for. Мы можем пройтись по каждому уровню, чтобы достать вложенные списки. Получается это таким образом (для двухуровневого):

names = []
for group in groups:
    for name in group:
        names.append(name)

У объекта list есть метод extend, который принимает на вход итерируемый объект и добавляет его элемент к вызывающему списку:

names = []
for group in groups:
    names.extend(group)

Подготовка данных для Data Mining на Python

Код курса
DPREP
Ближайшая дата курса
10 октября, 2022
Длительность обучения
32 ак.часов
Стоимость обучения
60 000 руб.

Так, мы визуально избавляемся от вложенных циклов, тем самым соблюдая принцип Python: Flat is better than nested. Вместо extend мы можем использовать обычное сложение: сумма двух списков дает один список. Поэтому код переписать можно так:

names = []
for group in groups:
    names += group

А ещё мы можем воспользоваться list comprehension, о котором мы говорили тут. Он позволяет немного подсократить код. Код на Python с использованием list comprehension выглядит так:

names = [name for group in groups
              for name in group]

Можно ли использовать * в list comprehension?

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

>>> numbers = [3, 4, 7]
>>> more_numbers = [2, 1, *numbers, 11, 18]
>>> more_numbers
[2, 1, 3, 4, 7, 11, 18]

Можем ли мы использовать данный оператор для распаковки итерируемых объектов внутри list comprehension? Ответ — нет.

>>> names = [*group for group in groups]
SyntaxError: iterable unpacking cannot be used in comprehension

Python был лишен данной возможности из-за проблем с читаемостью такого кода (см. PEP448).

Можно ли использовать Python-функцию sum?

Существует ещё трюк для разворачивания списка. И он не самый очевидный и подразумевает использование функции sum:

names = sum(groups, [])

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

Модуль itertools на все времена

Модуль itertools имеет целый набор различных классов и функций, которые предназначены для работы с итерируемыми объектами. Одним из таких классов является chain, который принимает на вход любое количество итерируемых аргументов и возвращает iterator. Этот iterator представляет собой цепочку переданных аргументов. Пример разворачивания списка Python с помощью chain:

>>> from itertools import chain
>>> chain(*groups)
<itertools.chain object at 0x7fc1b2d65bb0>
>>> list(chain(*groups))
['Hong', 'Ryan', 'Anthony', 'Wilhelmina', 'Margaret', 'Adrian']

Класс chain имеет специальный метод, который не требует предварительной распаковки:

list(chain.from_iterable(groups))

Мы рекомендуем использовать именно этот способ, поскольку он работает быстрее.

Больше подробностей о работе с различными структурами данных вы узнаете на наших образовательных курсах в лицензированном учебном центре обучения и повышения квалификации руководителей и ИТ-специалистов (менеджеров, архитекторов, инженеров, администраторов, Data Scientist’ов и аналитиков Big Data) в Москве:

Источники
  1. https://treyhunner.com/2021/11/how-to-flatten-a-list-in-python/
  2. https://docs.python.org/3/library/itertools.html#itertools.chain
Комментировать