Train-Test Split and Cross-Validation in Python: From Scratch to Scikit-Learn

Train-Test Split and Cross-Validation are statistical validation techniques used to evaluate how well machine learning models generalize to new, unseen data — preventing the model from merely memorizing the training dataset (overfitting).

Think of it like a teacher preparing students for a final exam: if the teacher only quizzes students on the exact same practice questions they memorized in class, high test scores do not prove they understand the subject. A fair evaluation holds back a set of surprise questions the students have never seen. Train-test split holds back one surprise exam, while cross-validation rotates through multiple surprise exams across the entire dataset to guarantee the score was not a fluke.

Why It Is Used: The Danger of Overfitting

A machine learning model can easily achieve 100% accuracy on its training data simply by memorizing noise and specific rows, while completely failing when deployed in production.

If you evaluate your model on the same data it was trained on, you have no way to distinguish true generalization from rote memorization. Data splitting isolates test records from the training process, providing an honest, unbiased measurement of predictive power.

The 3-Way Split: Train vs. Validation vs. Test Set

In professional machine learning workflows, a simple 2-way split is often insufficient when tuning hyperparameters. We use a 3-way split to prevent information leakage:

Which Validation Strategy Should You Pick?

The Mathematical Formulations

1. Train-Test Split Ratio

Given a dataset $D$ with $N$ total rows and test ratio $p$ (e.g. $p = 0.2$):

$$|D_{\text{train}}| = N \times (1 - p), \quad |D_{\text{test}}| = N \times p$$

2. K-Fold Cross-Validation Formulation

The dataset is partitioned into $k$ non-overlapping equal folds:

$$\text{Fold Size} = \left\lfloor \frac{N}{k} \right\rfloor$$

For each iteration $i \in \{1, 2, \dots, k\}$, one fold is held out for testing while the remaining $k-1$ folds train the model:

$$D_{\text{test}}^{(i)} = \text{Fold}_i, \quad D_{\text{train}}^{(i)} = D \setminus \text{Fold}_i$$

The overall expected performance ($\bar{s}$) and score standard deviation ($\sigma_s$) across all $k$ rounds are:

$\bar{s} = \frac{1}{k}\sum_{i=1}^{k} s_i, \quad \sigma_s = \sqrt{\frac{1}{k}\sum_{i=1}^{k}(s_i - \bar{s})^2}$

3. Stratified K-Fold Constraint

Maintains the class proportion of category $c$ across every fold:

$$\frac{\text{count}(c \in \text{Fold}_i)}{|\text{Fold}_i|} \approx \frac{\text{count}(c \in D)}{N} \quad \text{for every } i$$

Code: Step-by-Step Python Implementation

1. Setting Up the Dataset

import numpy as np
import pandas as pd

np.random.seed(42)

# Synthetic dataset: study hours predicting pass (1) or fail (0)
n_samples = 20
df = pd.DataFrame({
    "study_hours": np.random.uniform(0, 10, n_samples).round(1),
    "passed": np.random.choice([0, 1], size=n_samples, p=[0.4, 0.6])
})
print(df.head())

2. Train-Test Split — From Scratch

def train_test_split_scratch(data, test_size=0.2, random_state=None):
    if random_state is not None:
        np.random.seed(random_state)
    
    # 1. Shuffle row indices randomly
    shuffled_indices = np.random.permutation(len(data))
    
    # 2. Compute split boundary
    test_count = int(len(data) * test_size)
    
    test_idx = shuffled_indices[:test_count]
    train_idx = shuffled_indices[test_count:]
    
    return data.iloc[train_idx], data.iloc[test_idx]

train_df, test_df = train_test_split_scratch(df, test_size=0.2, random_state=42)
print("Training rows:", len(train_df))
print("Testing rows:", len(test_df))

3. Train-Test Split — Scikit-Learn (Production Version)

from sklearn.model_selection import train_test_split

X = df[["study_hours"]]
y = df["passed"]

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 shape:", X_train.shape)
print("X_test shape:", X_test.shape)

4. K-Fold Cross-Validation — From Scratch

from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

def k_fold_cross_validation_scratch(data, features, target, k=5, random_state=None):
    if random_state is not None:
        np.random.seed(random_state)
    
    indices = np.random.permutation(len(data))
    fold_size = len(data) // k
    scores = []
    
    for i in range(k):
        start = i * fold_size
        end = (i + 1) * fold_size if i != k - 1 else len(data)
        test_idx = indices[start:end]
        train_idx = np.setdiff1d(indices, test_idx)
        
        X_tr, y_tr = data.iloc[train_idx][features], data.iloc[train_idx][target]
        X_te, y_te = data.iloc[test_idx][features], data.iloc[test_idx][target]
        
        model = LogisticRegression()
        model.fit(X_tr, y_tr)
        preds = model.predict(X_te)
        scores.append(accuracy_score(y_te, preds))
        
    return np.array(scores)

cv_scores_manual = k_fold_cross_validation_scratch(
    df, features=["study_hours"], target="passed", k=5, random_state=42
)
print("Manual Fold Scores:", cv_scores_manual)
print(f"Mean Accuracy: {cv_scores_manual.mean():.2f} +/- {cv_scores_manual.std():.2f}")

5. Stratified K-Fold with Scikit-Learn Pipeline (Zero Leakage)

Best Practice: Encapsulate preprocessing scalers inside a Scikit-Learn Pipeline so features are fit strictly on the training folds:

from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline

# Pipeline guarantees scaling is fit ONLY on training folds
pipeline = make_pipeline(StandardScaler(), LogisticRegression())
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

scores = cross_val_score(pipeline, X, y, cv=skf, scoring="accuracy")
print("Stratified K-Fold Scores:", scores)
print(f"Cross-Validated Accuracy: {scores.mean():.2f} (+/- {scores.std():.2f})")

Common Pitfalls & How to Avoid Data Leakage

• Preprocessing Before Splitting: Never scale, impute, or encode the dataset before calling train_test_split or cross_val_score. Doing so leaks test set distribution parameters into training.
• Shuffling Time-Series Data: Randomly shuffling time-ordered datasets allows models to predict historical events using future observations. Always use TimeSeriesSplit for sequential data.
• Overfitting on the Test Set: Repeatedly modifying hyperparameters to maximize your test set score turns the test set into training data. Keep a separate Validation set or use Nested Cross-Validation.

Summary: The Quick Validation Checklist

• Use Train-Test Split for rapid prototyping and large-scale deep learning models.
• Use K-Fold Cross-Validation ($k=5$ or $10$) for reliable, low-variance evaluation on standard datasets.
• Use Stratified K-Fold whenever classification datasets contain imbalanced class distributions.
• Always combine preprocessing with Pipeline in cross-validation to guarantee zero data leakage.