AI Engineer Exhaustive
HardData Processing

Stratified Train/Test Split

Data Processing

Hard

Problem

Split features and labels while approximating the same class proportions in train and test. For each class with n_c samples, compute:

n_c^{\mathrm{test}} = \operatorname{round}(n_c\,t)

Here, t is test_size. When a class has more than one sample, cap its test count at $$ so at least one remains in training. Shuffle each class with np.random.default_rng(seed), combine the selected indices, sort each final index set, and return X_train, X_test, y_train, and y_test as NumPy arrays in a dictionary.

Theory

Stratified splitting divides a dataset into subsets while preserving the proportion of each class. If the original dataset has 70% class A and 30% class B, both the training and test sets will maintain approximately these same proportions. This is crucial for imbalanced classification problems.

Why Stratify?

Preserves class distribution: Random splitting can accidentally put most minority class samples in one split, leaving the other split unrepresentative.

Reliable evaluation: Test set performance reflects real-world distribution when class proportions are maintained.

Consistent training signal: Training set contains sufficient examples of all classes.

Essential for imbalanced data: When one class is rare (e.g., 1% fraud cases), random splitting could result in a test set with zero fraud cases.

The Problem with Random Splitting

Example scenario:

Possible random outcome:

This is lucky. But random splits can also produce:

The test set has double the minority class proportion - evaluation will be misleading.

Stratified Splitting Process

Step 1: Group samples by class label

Step 2: For each class, randomly split into train/test with the specified ratio

Step 3: Combine all class-specific train samples into the final training set

Step 4: Combine all class-specific test samples into the final test set

Result: Both sets have the same class proportions as the original data.

Mathematical Formulation

For a dataset with N samples and class distribution:

With train fraction f (e.g., 0.8 for 80% train):

Training set:

Test set:

Worked Example

Dataset: 100 samples

Split ratio: 80% train, 20% test

Stratified split calculation:

Class A:

Class B:

Class C:

Training set: 56 + 16 + 8 = 80 samples

Test set: 14 + 4 + 2 = 20 samples

Both sets preserve the original 70/20/10 distribution.

Handling Small Classes

When a class has very few samples, stratified splitting faces challenges:

Example: Class with 3 samples, 80/20 split

This works, but with only 2 samples, one split may have 2 or 0.

Solutions:

Stratified K-Fold Cross-Validation

Extends stratification to K-Fold:

Process:

  1. Group samples by class
  2. Within each class, divide into K folds
  3. Each fold contains proportional representation of all classes
  4. Iterate: use each fold as validation, others as training

Benefit: Every fold has representative class distribution for reliable cross-validation estimates.

Multi-Label Stratification

When samples can have multiple labels (multi-label classification):

Challenge: Simple stratification on single labels does not work

Iterative stratification algorithm:

  1. Order labels by frequency (rarest first)
  2. For each sample with the rarest label, assign to the fold with smallest proportion of that label
  3. Repeat for next rarest label

Goal: Approximately preserve the distribution of all label combinations

Implementation Considerations

Shuffling within classes: Randomly shuffle samples within each class before splitting to avoid ordering bias

Reproducibility: Set random seed for consistent splits across runs

Rounding: Floor function ensures train set gets the integer count; remaining go to test

Extremely rare classes: May need special handling or minimum sample requirements

Stratification vs Random Split

Use stratified split when:

Random split acceptable when:

Regression Stratification

For regression tasks, stratify on binned target values:

Process:

  1. Bin continuous target into discrete intervals
  2. Treat bins as pseudo-classes
  3. Apply stratified splitting on bins
  4. Ensures both splits have similar target distributions

Example: Income prediction

Where Stratified Splitting Shows Up

Examples

Example 1

Input
X = [0, 1, 2, 3, 4, 5], y = [0, 0, 0, 1, 1, 1], test_size = 0.33, seed = 42
Output
{"X_train": [0, 1, 4, 5], "X_test": [2, 3], "y_train": [0, 0, 1, 1], "y_test": [0, 1]}
Explanation
One seeded sample from each class enters the test split.

Example 2

Input
X = [[1, 0], [2, 0], [3, 0], [4, 0], [5, 0], [6, 0], [7, 0], [8, 0], [9, 0], [10, 0]], y = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1], test_size = 0.3, seed = 42
Output
{"X_train": [[1, 0], [2, 0], [5, 0], [6, 0], [7, 0], [8, 0], [9, 0]], "X_test": [[3, 0], [4, 0], [10, 0]], "y_train": [0, 0, 0, 0, 0, 1, 1], "y_test": [0, 0, 1]}

Hints

  1. Use np.flatnonzero(y == label) and rng.permutation for each class.
  2. Accumulate class indices, then apply np.sort before indexing X and y.

Requirements

Constraints

Starter Code

import numpy as np

def stratified_split(X: list, y: list, test_size: float = 0.2, seed: int = 42) -> dict:
    """
    Returns a dictionary with X_train, X_test, y_train, and y_test.
    """
    # Write code here
    pass

Test Cases

CaseMatches
Balanced binary—public
Imbalanced 2D features—public