Stratified Train/Val/Test Split - Data Preprocessing

Stratified Train/Val/Test Split

What you'll build

stratified_split(y, val_frac, test_frac, seed), which splits the sample indices 0 .. n-1 into train, validation and test sets class by class, so every split keeps the class proportions of the full dataset.

Background

A plain random split can leave a rare class badly under-represented, or missing entirely, from the validation or test set, which makes the metrics on those splits noisy or meaningless. Stratifying fixes the count of each class in each split. scikit-learn does this with train_test_split(..., stratify=y) (called twice for three splits); here the rounding rule is explicit so results are reproducible.

Algorithm

1. Create one generator before the loop: rng = np.random.default_rng(seed).
2. For each class c in ascending order (np.unique(y)):
   1. idx_c = np.where(y == c)[0] and perm = rng.permutation(idx_c); let n_c = len(idx_c).
   2. n_test = min(round_half_up(n_c * test_frac), n_c)
   3. n_val = min(round_half_up(n_c * val_frac), n_c - n_test) (validation gives way when the class is too small)
   4. The first n_test entries of perm go to test, the next n_val to val, and the rest to train.
3. Concatenate the per-class pieces and sort each split's indices in ascending order.

Rounding rule

Counts are rounded half up, per class: round_half_up(v) = int(np.floor(v + 0.5)), so 2.5 -> 3 and 0.5 -> 1. Python's built-in round uses banker's rounding (round(2.5) == 2), so don't use it here.

Inputs

• y: integer class labels, shape (n,)
• val_frac, test_frac: fractions in [0, 1] of each class to put in validation and test
• seed: int seed for np.random.default_rng

Expected Return Format

Return a dictionary with these keys:

• 'train': sorted int array of training indices
• 'val': sorted int array of validation indices
• 'test': sorted int array of test indices

Guarantees and conventions

• The three arrays are disjoint and together contain every index 0 .. n-1 exactly once.
• Each class's count in each split is fixed by the rule above; only which samples go where depends on the seed.
• The same seed always gives the same split. The exact indices depend on drawing the permutations in ascending class order from a single generator.

Constraints

• NumPy only; no sklearn.

Follow-up

• How would you stratify a regression target? (Bin it into quantiles first.)
• How do you split when several rows belong to the same user? (Group-aware splits, e.g. GroupKFold, so a user never appears in both train and test.)

Input: stratified_split(np.array([0]*10 + [1]*5), val_frac=0.2, test_frac=0.2, seed=0)
Output: {'train': [0, 1, 3, 5, 8, 9, 11, 13, 14], 'val': [2, 7, 12], 'test': [4, 6, 10]}

Input: stratified_split(np.array([0, 0, 0, 0, 0, 1, 1]), val_frac=0.1, test_frac=0.3, seed=0)
Output: {'train': [0, 1, 6], 'val': [3], 'test': [2, 4, 5]}

import numpy as np

def stratified_split(y, val_frac=0.15, test_frac=0.15, seed=0):
    """
    Stratified train/val/test split.

    Args:
        y: Integer class labels, shape (n,)
        val_frac: Fraction of each class to put in the validation set
        test_frac: Fraction of each class to put in the test set
        seed: Seed for np.random.default_rng

    Returns:
        dict with sorted int index arrays 'train', 'val', 'test'
    """
    # Your code here
    pass