Missing Data Imputation and Regression Pipeline

Machine Learning · Medium · Free problem

You are given a feature matrix X and a target vector y. Some entries of X are missing (represented as None). Build a small modeling pipeline that cleans the data and fits a linear model.

Implement a function with the signature:

```python def impute_and_fit(X, y): ... ```

  • X is an n x p matrix given as a list of rows; each entry is a float or None (missing).
  • y is the length-n target vector.

Requirements:

  1. Impute each column's missing (None) entries with that column's mean over the observed (non-None) values. If an entire column is missing, fill it with 0.0.
  2. Fit ordinary least squares WITH an intercept on the cleaned data.
  3. Return the coefficient vector [intercept, b1, b2, ...] as a plain list of floats (intercept first, then one slope per feature column).

Constraints: - Handle edge cases: all-None columns, single-feature inputs, and columns with several missing entries.

Example

```python impute_and_fit( [[1.0, None], [None, 3.0], [3.0, 5.0], [4.0, 6.0]], [2.1, 4.0, 6.2, 7.9], ) -> [-2.218474576271187, 1.7166101694915263, 0.5766101694915252] ```

Column 1's missing value is filled with the mean of its observed entries (1+3+4)/3, column 2's with (3+5+6)/3; then OLS with an intercept is fit and the returned list is [intercept, b_x1, b_x2].

Hints

  1. For interpolation, think about what happens at the edges of the series -- interpolate only fills interior NaNs because it needs two neighbors; you need a fallback for leading NaNs.
  2. The justification for MSE is not 'it is the default' -- it is that minimizing squared error is equivalent to maximizing the Gaussian log-likelihood, making OLS the MLE estimator under the standard noise assumption.
  3. For forward selection, maintain a selected list and a remaining list. At each round, loop over remaining, compute CV score for selected + [candidate], pick the best candidate, and move it from remaining to selected.

Worked Solution

How to Think About It: This is a two-stage pipeline: impute, then fit. The imputation rule is *mean over observed values*, computed per column (each column is filled independently with the average of its own non-missing entries; a fully missing column becomes all zeros). The fit is plain ordinary least squares with an intercept, and the required return value is the flat coefficient list [intercept, b1, b2, ...] -- intercept first, then one slope per feature.

Algorithm:

  • Impute: Load X into a float array, mapping None to NaN. For each column, take the mean of the observed (non-NaN) entries and write it into the missing slots. If every entry in a column is missing, set the whole column to 0.0.
  • Fit OLS with intercept: Prepend a column of ones to the imputed matrix and solve the least-squares system A @ beta = y. The first entry of beta is the intercept; the rest are the per-feature slopes.

Why MSE / OLS: Minimizing squared error is exactly maximum-likelihood estimation of beta under i.i.d. Gaussian noise y_i = x_i^T beta + eps_i: the Gaussian log-likelihood's beta-dependent term is -1/(2 sigma^2) * sum_i (y_i - x_i^T beta)^2, so maximizing likelihood is minimizing the sum of squared residuals. That is what np.linalg.lstsq returns.

Code:

```python import numpy as np

def impute_and_fit(X, y): # X is an n x p matrix (list of rows) of feature values; some entries are None # (missing). y is the length-n target. # Step 1: impute each column's missing entries with that column's MEAN over the # OBSERVED (non-None) values (if a whole column is missing, use 0.0). # Step 2: fit ordinary least squares WITH an intercept on the cleaned data. # Return the coefficient vector [intercept, b1, b2, ...] as a list of floats. n = len(X) p = len(X[0]) if n > 0 else 0

# Build design matrix with NaN for missing entries. Xf = np.array( [[np.nan if X[i][j] is None else float(X[i][j]) for j in range(p)] for i in range(n)], dtype=float, )

# Impute each column with the mean of its OBSERVED values (0.0 if all missing). for j in range(p): col = Xf[:, j] mask = np.isnan(col) if mask.all(): Xf[:, j] = 0.0 else: col[mask] = np.mean(col[~mask])

yv = np.asarray(y, dtype=float)

# OLS with intercept: prepend a column of ones and solve least squares. A = np.hstack([np.ones((n, 1)), Xf]) beta, *_ = np.linalg.lstsq(A, yv, rcond=None) return [float(b) for b in beta] ```

Complexity: Imputation is O(n * p). The least-squares solve on the n x (p+1) design matrix is O(n * p^2 + p^3). The returned list has length p + 1 (intercept followed by one coefficient per feature column).

Intuition

Linear interpolation for missing data implicitly assumes the underlying signal is smooth and the missingness is not informative -- i.e., the data is missing because of a sensor gap or a calendar gap, not because the value was extreme. In financial data this is often defensible for price series (a missing close price on a holiday can be linearly interpolated from adjacent days) but dangerous for event-driven data (a missing earnings report is not missing at random). Always ask why values are missing before choosing an imputation method.

The MSE-as-MLE argument is one of the most useful things to internalize in statistics: it tells you that OLS is not just a computational convenience but the statistically principled estimator when noise is Gaussian. The moment noise becomes heavy-tailed (as it often is in financial returns), MSE underweights outliers in the objective but overweights them in practice because a single large residual dominates. That is why practitioners use Huber loss or LAD regression for robustness. Forward selection is the simplest feature selection algorithm but suffers from a well-known flaw: it cannot remove a feature once added, so it can get stuck with an early bad choice. Regularization methods like Lasso are generally preferred in high-dimensional settings because they perform continuous feature shrinkage rather than discrete inclusion.

Related problems

Open the full interactive solver →