Splitting Data the Right Way

What you will learn in this lesson:

• What each of the train, validation, and test sets is actually for
• Why a plain random split can accidentally distort class proportions, and how stratified sampling fixes it
• Why shuffling before splitting matters for data that arrives in a meaningful order
• When a random split is the wrong tool, previewing the temporal leakage idea from the data leakage lesson

A Quick Reminder: Three Jobs, Three Splits

You'll meet the full mechanics of cross-validation later in this course, but the basic split into training, validation, and test data needs to be second nature before you evaluate a single model. As a quick reminder: the training set fits the model's parameters, the validation set is used to make decisions about the model (which hyperparameter, which algorithm), and the test set is touched exactly once, at the very end, to report an honest final number.

Why a Plain Random Split Can Distort Class Proportions

Suppose you have a dataset where 10% of customers churned and 90% didn't. A plain random 80/20 train/test split doesn't guarantee that exact 90/10 ratio survives the split. With enough bad luck, a random split could put a disproportionate share of the churned customers into the test set, leaving the training set looking more balanced than the real population actually is, or the reverse.

This matters most when the minority class is already small. Losing even a handful of positive examples to an unlucky split can meaningfully change both what the model learns and how reliable your evaluation metric is.

Stratified Sampling: Preserving the Proportions

Stratified sampling splits the data so that each split preserves the original class proportions. If the full dataset is 90% negative and 10% positive, a stratified 80/20 split guarantees both the training set and the test set are also approximately 90% negative and 10% positive, rather than leaving that to chance.

import numpy as np
from sklearn.model_selection import train_test_split

np.random.seed(0)
# A strongly imbalanced target: ~10% positive class
y = (np.random.rand(500) < 0.10).astype(int)
X = np.random.randn(500, 3)

# Plain random split: proportions can drift
X_train_r, X_test_r, y_train_r, y_test_r = train_test_split(X, y, test_size=0.2, random_state=1)
print("Random split    -> train positive rate: {:.3f}, test positive rate: {:.3f}".format(
    y_train_r.mean(), y_test_r.mean()))

# Stratified split: proportions preserved
X_train_s, X_test_s, y_train_s, y_test_s = train_test_split(
    X, y, test_size=0.2, random_state=1, stratify=y)
print("Stratified split -> train positive rate: {:.3f}, test positive rate: {:.3f}".format(
    y_train_s.mean(), y_test_s.mean()))

Run this a few times with different random_state values for the plain random split. You'll see its train/test positive rates drift further apart from each other, and from the true 10% rate, than the stratified version ever does. The stratify=y argument is the entire fix: pass it the target column, and scikit-learn handles the rest.

Use stratified splitting as the default for classification problems, especially with any class imbalance. Stratification has no equivalent for plain regression targets, since there are no discrete classes to balance, though scikit-learn does support binning a continuous target if you need an analogous guarantee there.

Why Shuffling Matters

train_test_split shuffles your data by default before splitting it. This matters because many real datasets arrive sorted, by date, by customer ID, or by whatever order they were collected in. Splitting an unshuffled, sorted dataset by simply taking the first 80% of rows as training data risks training on one systematic slice of the data (early customers, low ID numbers) and testing on a completely different slice, which can make your evaluation reflect that difference in the data rather than genuine model quality.

When Random Splitting Is the Wrong Tool

Shuffling is the right default specifically because most datasets don't have a meaningful time order baked into the prediction task. When they do, shuffling becomes actively harmful rather than protective: a model forecasting next month's sales should never be allowed to train on rows from after the period it's being tested on. The data leakage lesson covers this case, temporal leakage, in full.

The rule of thumb to carry forward from this lesson: shuffle and stratify by default for ordinary classification and regression, but switch to a chronological split the moment your data has a genuine time dimension tied to the prediction task.