Practice Now
pandas
dataframe
machine learning
stratified sampling
data splitting

Stratified splitting of pandas dataframe into training, validation and test set

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Machine LearningData Science

Introduction

Stratified splitting divides a dataset into training, validation, and test sets while preserving the class distribution from the original data. This is critical for classification tasks — especially with imbalanced classes — because random splitting can produce splits where minority classes are underrepresented or entirely absent. Scikit-learn's train_test_split with the stratify parameter handles this.

Why Stratified Splitting?

Consider a dataset with 95% class A and 5% class B. A random split might give a test set with 0% class B, making it impossible to evaluate the model on the minority class. Stratified splitting guarantees each split has approximately the same class proportions as the original.

python
1import pandas as pd
2from sklearn.model_selection import train_test_split
3
4df = pd.DataFrame({
5    'feature1': range(1000),
6    'feature2': range(1000),
7    'label': [0] * 950 + [1] * 50  # 95% class 0, 5% class 1
8})
9
10print(df['label'].value_counts(normalize=True))
11# 0    0.95
12# 1    0.05

Two-Step Split: Train / Validation / Test

Scikit-learn's train_test_split only splits into two sets. To get three sets, split twice:

python
1from sklearn.model_selection import train_test_split
2
3# Step 1: Split into train (60%) and temp (40%)
4X = df.drop('label', axis=1)
5y = df['label']
6
7X_train, X_temp, y_train, y_temp = train_test_split(
8    X, y,
9    test_size=0.4,
10    stratify=y,
11    random_state=42
12)
13
14# Step 2: Split temp into validation (20%) and test (20%)
15X_val, X_test, y_val, y_test = train_test_split(
16    X_temp, y_temp,
17    test_size=0.5,       # 50% of 40% = 20% of total
18    stratify=y_temp,
19    random_state=42
20)
21
22print(f"Train: {len(X_train)}, Val: {len(X_val)}, Test: {len(X_test)}")
23# Train: 600, Val: 200, Test: 200

Verify the class distribution is preserved:

python
1print(f"Original:   {y.value_counts(normalize=True).to_dict()}")
2print(f"Train:      {y_train.value_counts(normalize=True).to_dict()}")
3print(f"Validation: {y_val.value_counts(normalize=True).to_dict()}")
4print(f"Test:       {y_test.value_counts(normalize=True).to_dict()}")
5# All should show approximately {0: 0.95, 1: 0.05}

Working Directly with DataFrames

If you prefer to keep the DataFrame intact rather than separating X and y:

python
1# Split the entire DataFrame, stratify by the label column
2train_df, temp_df = train_test_split(
3    df, test_size=0.4, stratify=df['label'], random_state=42
4)
5
6val_df, test_df = train_test_split(
7    temp_df, test_size=0.5, stratify=temp_df['label'], random_state=42
8)
9
10print(train_df.shape, val_df.shape, test_df.shape)
11# (600, 3) (200, 3) (200, 3)

Custom Split Ratios

For a 70/15/15 split:

python
1# 70% train, 30% temp
2X_train, X_temp, y_train, y_temp = train_test_split(
3    X, y, test_size=0.3, stratify=y, random_state=42
4)
5
6# 50% of 30% = 15% val, 15% test
7X_val, X_test, y_val, y_test = train_test_split(
8    X_temp, y_temp, test_size=0.5, stratify=y_temp, random_state=42
9)

For an 80/10/10 split:

python
1X_train, X_temp, y_train, y_temp = train_test_split(
2    X, y, test_size=0.2, stratify=y, random_state=42
3)
4
5X_val, X_test, y_val, y_test = train_test_split(
6    X_temp, y_temp, test_size=0.5, stratify=y_temp, random_state=42
7)

Multi-Class Stratified Splitting

Stratification works the same for multi-class problems:

python
1df = pd.DataFrame({
2    'text': [f'sample_{i}' for i in range(500)],
3    'category': ['sports'] * 200 + ['politics'] * 150 + ['tech'] * 100 + ['health'] * 50
4})
5
6train_df, temp_df = train_test_split(
7    df, test_size=0.3, stratify=df['category'], random_state=42
8)
9
10val_df, test_df = train_test_split(
11    temp_df, test_size=0.5, stratify=temp_df['category'], random_state=42
12)
13
14# Verify distribution
15for name, split in [('Original', df), ('Train', train_df), ('Val', val_df), ('Test', test_df)]:
16    dist = split['category'].value_counts(normalize=True).round(3)
17    print(f"{name}: {dist.to_dict()}")

Alternative: StratifiedShuffleSplit

For repeated stratified splits (useful in cross-validation):

python
1from sklearn.model_selection import StratifiedShuffleSplit
2
3sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
4
5for train_idx, test_idx in sss.split(X, y):
6    X_train, X_test = X.iloc[train_idx], X.iloc[test_idx]
7    y_train, y_test = y.iloc[train_idx], y.iloc[test_idx]

Stratified K-Fold for Cross-Validation

When you do not want a fixed split and prefer cross-validation:

python
1from sklearn.model_selection import StratifiedKFold
2
3skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
4
5for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)):
6    X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
7    y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]
8    print(f"Fold {fold}: train={len(X_train)}, val={len(X_val)}")

Common Pitfalls

  • Too few samples per class: Stratification fails if a class has fewer samples than the number of splits. With very rare classes, consider merging categories or using StratifiedShuffleSplit which is more forgiving.
  • Data leakage: Split before any preprocessing (normalization, feature engineering). Fitting a scaler on the full dataset before splitting leaks test set statistics into training.
  • Forgetting random_state: Without a fixed random_state, each run produces different splits, making experiments irreproducible. Always set it for reproducibility.
  • Ignoring group structure: If multiple rows belong to the same entity (e.g., multiple visits from one patient), use GroupShuffleSplit or StratifiedGroupKFold to keep all rows from one group in the same split.
  • Continuous targets: stratify only works with categorical labels. For regression, use binning (pd.cut()) to create strata, or use train_test_split without stratification.

Summary

  • Use train_test_split with stratify=y to preserve class distribution across splits
  • Split twice for train/val/test: first 60/40, then split the 40% into 50/50
  • Always verify class proportions in each split with value_counts(normalize=True)
  • Use StratifiedKFold for cross-validation with stratification
  • Set random_state for reproducibility and split before any preprocessing
Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.