Шесть способов сделать flatten в Python — какой быстрее в 500 раз

Если в коде встречается список списков и нужно превратить его в один плоский, на ум приходит как минимум пять подходов — от обычного for до itertools.chain и NumPy. Какие из них работают за миллисекунды, а какие — за секунды? Перевод и адаптация туториала Real Python с разбором всех способов и замером производительности.

Flatten — операция приведения вложенной структуры к одномерной. Самый распространённый кейс: матрица представлена как список списков, а вам нужен плоский список всех значений, чтобы передать в алгоритм или ML-модель. В отличие от других языков, в Python нет встроенного .flatten() у списков — поэтому стандартное «как это сделать?» имеет несколько ответов разной степени читаемости и эффективности.

Пусть у нас есть матрица 4×4:

			matrix = [
    [9, 3, 8, 3],
    [4, 5, 2, 8],
    [6, 4, 3, 1],
    [1, 0, 4, 5],
]
		

Цель — получить:

			[9, 3, 8, 3, 4, 5, 2, 8, 6, 4, 3, 1, 1, 0, 4, 5]
		

Способ 1. Цикл for + .extend()

Самый прямолинейный и читаемый вариант: проходим по каждой подсписку и добавляем её элементы к итоговому списку через .extend().

			def flatten_extend(matrix):
    flat_list = []
    for row in matrix:
        flat_list.extend(row)
    return flat_list
		

.extend() принимает любой iterable и поэлементно добавляет его в конец списка. Альтернативно можно использовать augmented-concatenation +=:

			def flatten_concat(matrix):
    flat_list = []
    for row in matrix:
        flat_list += row
    return flat_list
		

Обе функции делают одно и то же. Real Python отмечает: flatten_extend читается чуть лучше, потому что метод named — намекает на смысл операции. По скорости варианты идут в первой тройке.

Способ 2. List comprehension

Pythonic-альтернатива однострочнику. Двойной for внутри comprehension перебирает сначала строки, потом элементы строк:

			def flatten_comp(matrix):
    return [item for row in matrix for item in row]
		

Подход компактнее цикла и не плодит промежуточных переменных. Скорость — заметно медленнее .extend() из-за двух вложенных циклов на уровне comprehension, но всё ещё в диапазоне десятков миллисекунд.

Способ 3. itertools.chain() — экономно по памяти

chain() возвращает итератор, а не список. Это значит, что вы можете обходить элементы потокового, не материализуя весь массив в памяти — полезно, когда матрица большая и хранить полный flat-список нерационально.

			from itertools import chain

def flatten_chain(matrix):
    return list(chain.from_iterable(matrix))
		

На большой матрице chain.from_iterable примерно в три раза медленнее цикла .extend(): накладные расходы на материализацию через list(). Но если итоговый список как таковой не нужен — можно отдать сам итератор потребителю и сэкономить и время, и память.

Способ 4. functools.reduce() и sum() — почему так медленно

Эти варианты часто появляются в туториалах для красоты, но на практике плохо масштабируются. Причина одна: каждая итерация создаёт новый промежуточный список вместо мутации существующего.

			from functools import reduce
from operator import iconcat

# Быстрая версия — c iconcat (in-place):
def flatten_reduce_iconcat(matrix):
    return reduce(iconcat, matrix, [])

# Медленная версия — с lambda и созданием новых списков:
def flatten_reduce_lambda(matrix):
    return reduce(lambda x, y: x + y, matrix)
		

iconcat мутирует аккумулятор и держится в топ-3 по скорости. А вот вариант с lambda и оператором + на каждом шаге копирует обе стороны — отсюда и взрывной рост времени на больших данных.

С sum() ситуация ровно такая же:

			def flatten_sum(matrix):
    return sum(matrix, [])
		

sum() работает на любых сложениях, но для списков это О(n²) по аллокациям. На матрице 1000×1000 он в ~500 раз медленнее цикла с .extend().

Способ 5. Произвольная вложенность — рекурсия и стек

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

			def flatten_recursive(items):
    flat = []
    for x in items:
        if isinstance(x, list):
            flat.extend(flatten_recursive(x))
        else:
            flat.append(x)
    return flat
		

Рекурсия удобна, пока глубина не упирается в sys.getrecursionlimit() (по умолчанию 1000). Для очень глубоких структур стоит переписать через итеративный обход со стеком:

			def flatten_stack(items):
    flat = []
    stack = [iter(items)]
    while stack:
        try:
            x = next(stack[-1])
        except StopIteration:
            stack.pop()
            continue
        if isinstance(x, list):
            stack.append(iter(x))
        else:
            flat.append(x)
    return flat
		

Способ 6. NumPy для data-science

Если вы работаете с матрицами в data-science, у numpy.ndarray есть готовый метод .flatten(). Он быстрее любого варианта на чистых Python-списках, потому что под капотом работает с непрерывным блоком памяти.

			import numpy as np

matrix = np.array([
    [9, 3, 8, 3],
    [4, 5, 2, 8],
    [6, 4, 3, 1],
    [1, 0, 4, 5],
])
matrix.flatten()
# array([9, 3, 8, 3, 4, 5, 2, 8, 6, 4, 3, 1, 1, 0, 4, 5])
		

Замер производительности

Real Python прогнали все варианты на матрице 1000×1000 через timeit. Результаты, отсортированные по времени:

			flatten_concatenation()..........5.29 ms
flatten_reduce_iconcat().........5.78 ms
flatten_extend()................. 6.62 ms
flatten_chain()................. 15.63 ms
flatten_comprehension()..........23.69 ms
flatten_reduce_lambda().......2645.11 ms
flatten_reduce_concat().......2652.81 ms
flatten_reduce_add()..........2658.96 ms
flatten_sum()................. 2684.48 ms
		

Между первыми тремя и последними четырьмя — пропасть в 400–500 раз. Причина: топовые функции мутируют существующий список, а reduce(+), reduce(operator.add), reduce(lambda x, y: x + y) и sum() на каждой итерации создают копии.

FAQ

Выводы

Если вам нужен один-в-один поведенческий рецепт: for row in matrix: flat.extend(row). Этот код понятен любому, кто знает Python, и упирается в производительность только когда матрица переваливает за миллионы элементов. Если объёмы такие, что нужна потоковая обработка — переключитесь на itertools.chain. Если данные numerical — используйте NumPy.

Главное, чего стоит избегать — sum(matrix, []) и reduce с + или lambda. Они смотрятся изящно в туториалах, но превращают O(n) в O(n²) и ставят вас в первую же боттлнек-историю на проде.

Оригинальный туториал с пошаговыми объяснениями и сравнением — на Real Python, автор — Leodanis Pozo Ramos.