Что такое стратификация (stratification) при разбиении данных?

«Что такое стратификация (stratification) при разбиении данных?» — вопрос из категории Предобработка данных, который задают на 26% собеседований Data Scientist / ML Инженер. Ниже — развёрнутый ответ с разбором ключевых моментов.

Ответ

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

Зачем это нужно? Без стратификации при случайном разбиении в тестовой выборке может оказаться непропорционально мало примеров редкого класса, что приведет к некорректной оценке модели, особенно для задач с несбалансированными классами.

Пример на Python с scikit-learn:

from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np

# Создадим синтетический несбалансированный датасет
# 1000 примеров, целевая переменная 'target' имеет классы 0 и 1 (90% и 10%)
data = pd.DataFrame({'feature': np.random.randn(1000)})
data['target'] = np.random.choice([0, 1], size=1000, p=[0.9, 0.1])

print('Исходное распределение классов:')
print(data['target'].value_counts(normalize=True))
# 0: 0.90, 1: 0.10

# Разбиение БЕЗ стратификации
X_train_naive, X_test_naive, y_train_naive, y_test_naive = train_test_split(
    data[['feature']], data['target'], test_size=0.2, random_state=42
)
print('nБез стратификации в тесте:')
print(y_test_naive.value_counts(normalize=True))
# Может получиться, например, 0: 0.92, 1: 0.08 — класс 1 представлен хуже.

# Разбиение СО стратификацией
X_train, X_test, y_train, y_test = train_test_split(
    data[['feature']], data['target'],
    test_size=0.2,
    stratify=data['target'],  # Ключевой параметр
    random_state=42
)
print('nСо стратификацией в тесте:')
print(y_test.value_counts(normalize=True))
# Гарантированно будет ~0: 0.90, 1: 0.10 — распределение сохранено.

Где еще применяется?

  • StratifiedKFold / StratifiedShuffleSplit: Для кросс-валидации, чтобы в каждом фолде было то же распределение классов, что и в исходных данных.
  • Стратификация по нескольким признакам: Можно создать составной стратифицирующий признак.

Ограничение: Применима только к категориальным (дискретным) целевым переменным.