Группировка списков Python

Группировка списка означает разбиение его элементов на подколлекции, объединённые общим ключом — например, группировка слов по первой букве или группировка записей по полю категории. Python предлагает три основных подхода: ручной цикл с collections.defaultdict, itertools.groupby из стандартной библиотеки и словарные включения. В этой главе объясняется каждый метод, когда выбрать тот или иной, а также рассматриваются типичные ошибки.

Связанные главы: Списки Python · Методы списков · Включения списков · Перебор списков · Модуль collections

Что такое «группировка»

Дан плоский список и функция ключа, которая отображает каждый элемент на метку группы. Цель — получить отображение от каждой метки к списку принадлежащих ей элементов:

['apple', 'banana', 'avocado', 'blueberry', 'cherry', 'apricot']
  key = first letter
  →  {'a': ['apple', 'avocado', 'apricot'],
      'b': ['banana', 'blueberry'],
      'c': ['cherry']}

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

Метод 1: ручной цикл с defaultdict

collections.defaultdict — наиболее распространённый и гибкий подход. При обращении к несуществующему ключу defaultdict(list) автоматически создаёт для него пустой список, поэтому проверка if key in d не нужна.

from collections import defaultdict

words = ['apple', 'banana', 'avocado', 'blueberry', 'cherry', 'apricot']

by_letter = defaultdict(list)
for word in words:
    by_letter[word[0]].append(word)

for letter, group in sorted(by_letter.items()):
    print(f'{letter}: {group}')
# a: ['apple', 'avocado', 'apricot']
# b: ['banana', 'blueberry']
# c: ['cherry']

Зачем использовать defaultdict вместо обычного словаря?

С обычным dict нужна явная проверка перед первым append:

# Plain dict — more boilerplate, same result
by_letter = {}
for word in words:
    if word[0] not in by_letter:
        by_letter[word[0]] = []
    by_letter[word[0]].append(word)

Более краткой альтернативой является dict.setdefault:

by_letter = {}
for word in words:
    by_letter.setdefault(word[0], []).append(word)

setdefault подходит для коротких скриптов, но defaultdict быстрее (нет повторных поисков ключей) и явнее выражает намерение.

Группировка по вычисляемому ключу

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

from collections import defaultdict

numbers = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]

by_parity = defaultdict(list)
for n in numbers:
    by_parity['even' if n % 2 == 0 else 'odd'].append(n)

print('even:', sorted(by_parity['even']))  # even: [2, 4, 6]
print('odd:', sorted(by_parity['odd']))    # odd: [1, 1, 3, 3, 5, 5, 5, 9]

Группировка списка словарей

Это наиболее распространённый практический сценарий — группировка строк данных по значению поля:

from collections import defaultdict

data = [
    {'category': 'fruit', 'name': 'apple'},
    {'category': 'vegetable', 'name': 'carrot'},
    {'category': 'fruit', 'name': 'banana'},
    {'category': 'vegetable', 'name': 'broccoli'},
]

grouped = defaultdict(list)
for item in data:
    grouped[item['category']].append(item['name'])

for category, names in grouped.items():
    print(f'{category}: {names}')
# fruit: ['apple', 'banana']
# vegetable: ['carrot', 'broccoli']

Метод 2: itertools.groupby

itertools.groupby группирует последовательные элементы с одинаковым ключом. Это полезно, когда нужно сохранить структуру серий или когда данные уже отсортированы и нежелательно строить весь словарь сразу (функция работает лениво, в потоковом режиме).

from itertools import groupby

words = ['apple', 'banana', 'avocado', 'blueberry', 'cherry', 'apricot']

# groupby only groups consecutive elements, so sort first
words_sorted = sorted(words, key=lambda w: w[0])

for letter, group in groupby(words_sorted, key=lambda w: w[0]):
    print(f'{letter}: {list(group)}')
# a: ['apple', 'avocado', 'apricot']
# b: ['banana', 'blueberry']
# c: ['cherry']

Критическая ловушка: сортировка перед groupby

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

from itertools import groupby

# Unsorted input — groupby produces WRONG results
numbers = [1, 1, 2, 3, 3, 1, 2, 2]
for key, group in groupby(numbers):
    print(f'{key}: {list(group)}')
# 1: [1, 1]   ← first run of 1s
# 2: [2]
# 3: [3, 3]
# 1: [1]      ← second run of 1s — NOT merged with the first!
# 2: [2, 2]

Всегда сортируйте по той же функции ключа перед вызовом groupby:

numbers_sorted = sorted(numbers)
for key, group in groupby(numbers_sorted):
    print(f'{key}: {list(group)}')
# 1: [1, 1, 1]
# 2: [2, 2, 2]
# 3: [3, 3]

Когда groupby особенно полезен: потоковая обработка больших данных

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

from itertools import groupby

# Grouping namedtuple records
from collections import namedtuple

Product = namedtuple('Product', ['category', 'name', 'price'])
products = [
    Product('dairy', 'milk', 1.10),
    Product('fruit', 'apple', 1.20),
    Product('fruit', 'banana', 0.50),
    Product('vegetable', 'broccoli', 1.50),
    Product('vegetable', 'carrot', 0.80),
]
# products is already sorted by category here

for category, group in groupby(products, key=lambda p: p.category):
    items = list(group)
    print(f'{category}: {[p.name for p in items]}')
# dairy: ['milk']
# fruit: ['apple', 'banana']
# vegetable: ['broccoli', 'carrot']

Метод 3: словарное включение

Словарное включение строит сгруппированный словарь в одном выражении. Это компактно, но имеет недостаток: внутреннее включение списка повторно просматривает весь ввод для каждого уникального ключа, что даёт сложность O(n × k), где k — количество уникальных ключей. Для небольших списков это приемлемо; для больших предпочтительнее defaultdict.

words = ['apple', 'banana', 'avocado', 'blueberry', 'cherry', 'apricot']

# Collect unique keys first, then build each group
letters = sorted(set(w[0] for w in words))
grouped = {letter: [w for w in words if w[0] == letter] for letter in letters}

for letter, group in grouped.items():
    print(f'{letter}: {group}')
# a: ['apple', 'avocado', 'apricot']
# b: ['banana', 'blueberry']
# c: ['cherry']

Этот метод наиболее читаем, когда набор ключей мал и уже известен — например, при группировке результатов True/False или фиксированного набора категорий.

Агрегация групп после группировки

Распространённым продолжением группировки является агрегация: вычисление суммы, среднего значения, минимума или количества по группе. Совместите defaultdict(list) со стандартной арифметикой Python:

from collections import defaultdict

scores = [
    ('Alice', 90), ('Bob', 75), ('Alice', 85),
    ('Bob', 88), ('Carol', 92),
]

by_student = defaultdict(list)
for name, score in scores:
    by_student[name].append(score)

for student, student_scores in sorted(by_student.items()):
    avg = sum(student_scores) / len(student_scores)
    print(f'{student}: scores={student_scores}, avg={avg:.1f}')
# Alice: scores=[90, 85], avg=87.5
# Bob: scores=[75, 88], avg=81.5
# Carol: scores=[92], avg=92.0

Выбор подходящего метода

Типичные ошибки

Забыть отсортировать данные перед groupby. groupby объединяет только последовательные одинаковые ключи. Всегда применяйте sorted() по той же функции ключа перед передачей данных в groupby.

Не сохранять list(group) сразу. Итератор группы из groupby исчерпывается, как только внешний цикл for переходит к следующему ключу. Преобразуйте его в список внутри тела цикла, если он понадобится более одного раза.

Изменять входной список во время группировки. Добавление или удаление элементов из списка во время цикла группировки приводит к непредсказуемым результатам. Сначала постройте сгруппированный словарь, а затем изменяйте элементы.

defaultdict в repr выглядит иначе. defaultdict(list, {...}) отображается иначе, чем обычный dict в repr. Оберните его в dict(grouped), если нужен вывод в виде обычного словаря.

Практика