Before training any model, you split your dataset so some data is held back to honestly evaluate the model on examples it never saw during training. In scikit-learn, this is one function call — but a few of its parameters matter more than they look.
The Basic Call
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(
X, y,
test_size=0.2, # 20% of data held out for testing
random_state=42 # makes the split reproducible
)
print(X_train.shape, X_test.shape) # e.g. (800, 5) (200, 5) for a 1000-row datasetWhy random_state Matters
train_test_split shuffles the data before splitting. Without a fixed random_state, you get a different train/test split every run — meaning your model's reported accuracy can shift slightly each time, purely from which rows happened to land in the test set. Fixing random_state makes results reproducible and comparisons between model versions fair.
stratify — Critical for Classification
# Without stratify: class proportions in train/test can differ from the original data
# With stratify: train and test both keep roughly the SAME class proportions
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)On an imbalanced dataset (say, 5% fraud cases), a random split without stratify could easily put almost no fraud examples in the test set — making evaluation metrics unreliable. Always use stratify=y for classification problems.
Three-Way Split — Train / Validation / Test
# First split off the test set
X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.15, random_state=42, stratify=y)
# Then split the remainder into train and validation
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.176, random_state=42, stratify=y_temp)
# 0.176 of the remaining 85% ≈ 15% of the original data -> final split is roughly 70/15/15The validation set is used to tune hyperparameters and compare models; the test set is touched exactly once, at the very end, to report final performance. See Cross-Validation for a more robust alternative to a single validation split.
Common Mistakes
- Preprocessing before splitting — scaling, imputing, or feature-selecting on the full dataset before splitting leaks test-set information into training. Always split first, then fit preprocessing only on the training set. See Data Leakage.
- Forgetting
stratify=yon an imbalanced classification dataset. - Touching the test set more than once — repeatedly checking test performance while tuning turns it into a second validation set, and your final reported number becomes overly optimistic.
Interview Relevance
Q: "Why should you fit a scaler on the training set only, not the full dataset before splitting?" Because fitting on the full dataset lets statistics from the test set (mean, std, min/max) influence the transformation applied to training data — the model implicitly "sees" test-set information before it should, inflating evaluation metrics.
Practice Question
You're building a fraud classifier where only 2% of transactions are fraudulent. Write the train_test_split call that correctly handles this class imbalance.
