Ravindra BagaleCourses & study guides मराठी Track your progress

31. Train Test Split in Machine Learning

31.3 Stratify and random_state Explained Simply

Why. Two small arguments save you from silent bias: stratify and random_state.

What.

  • random_state — fixes the shuffle so the same split repeats. Good for teaching and debugging.
  • stratify=y — keeps the same yes/no ratio in train and test. Useful when late orders are rare (e.g. only 10% late).

How.

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.DataFrame({
 "x": range(20),
 "was_late": [0]*18 + [1, 1], # rare late cases (fictional)
})
X = df[["x"]]
y = df["was_late"]

X_tr, X_te, y_tr, y_te = train_test_split(
 X, y, test_size=0.25, random_state=42, stratify=y
)
print("train late rate", y_tr.mean())
print("test late rate", y_te.mean())

What you should see. Train and test late rates stay close. Without stratify, a tiny dataset might put both late rows only in train — then test never sees a late case.

Ravindra Bagale's Tip

Classification + rare class asel tar stratify almost always. Regression (number predict) madhe stratify usually nahi. Fark lakshat theva.