Train-Test Split Concepts

Why split data?

When you evaluate a model, you want to test on data it hasn’t seen.

• Train set: used to learn patterns
• Test set: used only for final evaluation

This simulates real-world performance.

The biggest danger: data leakage

Leakage happens when information from the test set influences training.

Examples:

• Scaling using mean/std computed on the full dataset
• Filling missing values using overall mean (including test)
• Feature engineering that uses future information

Basic split with scikit-learn

train_test_split

import pandas as pd
from sklearn.model_selection import train_test_split
 
X = pd.DataFrame({"age": [20, 21, 22, 23, 24], "score": [80, 85, 78, 90, 88]})
y = pd.Series([0, 0, 0, 1, 1])
 
X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    random_state=42,
    stratify=y,
)
 
print(X_train)
print(X_test)

Stratification

If your target classes are imbalanced, use stratify=y so train/test have similar class distribution.

Time-series splits

For time series, you often do not shuffle. You train on past and test on future.

Good practice

• Keep a final test set untouched.
• Use cross-validation on training data for tuning.
• Put preprocessing inside a pipeline.

Visualize it

diagram Train/test split mermaid

  flowchart LR
  A["Full dataset"] --> B["train_test_split(test_size=0.2)"]
  B --> C["Train set (80%)"]
  B --> D["Test set (20%, untouched)"]
  C --> E["Fit model"]
  E --> F["Evaluate on D"]

The full dataset is split once; the model learns only from train and is evaluated only on the untouched test set.

sketch Splitting one dataset into train and test p5.js Pause Restart

The same rows, rearranged into an 80/20 split — train teaches the model, test grades it.

From the book

Splitting is really just random sampling. McKinney shows df.sample(n=...) for sampling without replacement and numpy.random.permutation(len(df)) for shuffling row order — train_test_split’s random_state argument does the same kind of seeded shuffle under the hood, which is why the same random_state always reproduces the same split.

🧪 Try It Yourself

Exercise 1 – Basic train/test split

pch.exerciseTag pch.exerciseTitle

Exercise 2 – Keep class balance with stratify

pch.exerciseTag pch.exerciseTitle

Exercise 3 – Reproducible splits with random_state

pch.exerciseTag pch.exerciseTitle

Next

Once your split strategy is settled, wire scaling, encoding, and the model together with a Preprocessing Pipeline (scikit-learn).