Leakage-Safe Imputation and Scaling - Data Preprocessing

Leakage-Safe Imputation and Scaling

What you'll build

fit_transform_train_test(X_train, X_test, strategy): fill missing values and standardise the features, learning every statistic from the training split only and applying the same transformation to both splits.

Background

A classic preprocessing bug is to compute imputation values or scaling statistics on the whole dataset before splitting. The test set then leaks into training: its values shift the medians, means and standard deviations the model is trained with, and the offline score becomes optimistic. The fix is the scikit-learn Pipeline pattern: fit on train, transform train and test. This problem is SimpleImputer(strategy=...) followed by StandardScaler().

Algorithm

1. Fit imputation on train: for each column j, fill[j] is the median (strategy='median') or the mean (strategy='mean') of the observed (non-NaN) training values. The median follows NumPy: the average of the two middle values for an even count. If a column has no observed training values, fill[j] = 0.0.
2. Impute both splits with the same fill values.
3. Fit scaling on the imputed train split: mean[j] and std[j], the population standard deviation (ddof=0).
4. Standardise both splits: (X - mean) / scale with scale[j] = std[j], or 1 when std[j] == 0.

Inputs

• X_train: float array (n_train, n_features), may contain np.nan
• X_test: float array (n_test, n_features), may contain np.nan
• strategy: 'median' (default) or 'mean'

Expected Return Format

Return a dictionary with these keys:

• Transformed data: 'X_train', 'X_test': float arrays with the input shapes and no NaNs
• Fitted statistics, each of shape (n_features,) and computed from the training split only:
  • 'fill': the imputation values
  • 'mean': column means of the imputed training data
  • 'std': column population stds of the imputed training data (the raw value, 0.0 included)

Edge cases and conventions

• Zero-variance column (std[j] == 0): divide by 1 instead, so it becomes 0 on train, as StandardScaler does.
• All-NaN training column: fill[j] = 0.0.
• Don't modify the inputs: work on copies.
• Nothing in X_test may influence fill, mean, std or the transformed X_train.

Constraints

• NumPy only; no sklearn. A loop over columns is fine for the fill values; vectorise the rest.

Follow-up

• With k-fold cross-validation, where must the imputer and scaler be fitted? (Inside each fold, on that fold's training part.)
• Why might a "was missing" indicator column be worth adding before imputing?

Input: fit_transform_train_test(np.array([[1.0, 10.0], [2.0, np.nan], [3.0, 30.0], [np.nan, 40.0]]), np.array([[100.0, np.nan], [np.nan, 0.0]]))
Output: {'fill': [2.0, 30.0], 'mean': [2.0, 27.5], 'std': [0.7071, 10.8972], 'X_test': [[138.5929, 0.2294], [0.0, -2.5236]], 'X_train': [[-1.4142, -1.6059], [0.0, 0.2294], [1.4142, 0.2294], [0.0, 1.1471]]}

Input: fit_transform_train_test(np.array([[5.0, 1.0], [5.0, 2.0], [np.nan, 3.0]]), np.array([[7.0, 2.0]]))
Output: {'fill': [5.0, 2.0], 'std': [0.0, 0.8165], 'X_train'[:, 0]: [0.0, 0.0, 0.0], 'X_test'[0, 0]: 2.0, ...}

import numpy as np

def fit_transform_train_test(X_train, X_test, strategy='median'):
    """
    Impute missing values and standardise, fitting statistics on the training split only.

    Args:
        X_train: Training features with NaNs, array (n_train, n_features)
        X_test: Test features with NaNs, array (n_test, n_features)
        strategy: 'median' or 'mean' imputation

    Returns:
        dict with 'X_train', 'X_test' (transformed arrays) and the fitted
        statistics 'fill', 'mean', 'std' (each shape (n_features,))
    """
    # Your code here
    pass