Code
machine-learning
import numpy as np
from sklearn.model_selection import train_test_split, StratifiedShuffleSplit
X = np.random.rand(100, 5)
y = np.random.randint(0, 2, 100)
# Standard split
X_tr, X_te, y_tr, y_te = train_test_split(
X, y, test_size=0.2, random_state=42, shuffle=True)
# Stratified split preserves class proportions
X_tr, X_te, y_tr, y_te = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y)
# Three-way split: train/val/test
X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.25, random_state=42)
# Repeated stratified splits
sss = StratifiedShuffleSplit(n_splits=5, test_size=0.2, random_state=42)
for train_idx, test_idx in sss.split(X, y):
print(len(train_idx), len(test_idx))