Train-Test Split in Python

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 dataset

Why 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/15

The 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=y on 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.

Related ML Notes

• Train-Test Split (Concept)
• Data Leakage
• Cross-Validation