Handle Missing Data (Mean Imputation) - Data Preprocessing

Handle Missing Data (Mean Imputation)

What you'll build

A function that replaces every missing value (NaN) in a 1-D array with the mean of the values that are present.

Background

Most models can't consume NaN. The simplest fix is imputation: fill each gap with a statistic of the observed values. Mean imputation keeps the column mean unchanged; it is SimpleImputer(strategy='mean') in scikit-learn. (In a real pipeline the mean must come from the training split only: see Leakage-Safe Imputation and Scaling.)

Formula

observed = { x[i] : x[i] is not NaN }
fill     = sum(observed) / |observed|
out[i]   = fill if x[i] is NaN else x[i]

Inputs

• arr: 1-D float array of shape (n,) with at least one non-NaN value; missing entries are np.nan (tests may pass a Python list)

Expected return format

A new float array of shape (n,) with no NaNs. Values that were present are unchanged.

Edge cases and conventions

• An array with no NaNs is returned unchanged (as a copy).
• Don't modify the input: arr[mask] = v on the caller's array would change their data. Work on a copy.
• np.mean of an array containing NaN is NaN, so you must exclude the missing values when computing the fill.

Constraints

• Use NumPy; vectorise (no Python loop over elements).

Follow-up

• When is the median a better fill than the mean?
• Why can adding a binary "was missing" feature help a model?

Input: fill_missing_with_mean(np.array([1.0, np.nan, 3.0, np.nan, 5.0]))
Output: [1.0, 3.0, 3.0, 3.0, 5.0]

Input: fill_missing_with_mean(np.array([np.nan, 10.0, 20.0]))
Output: [15.0, 10.0, 20.0]

import numpy as np

def fill_missing_with_mean(arr: np.ndarray) -> np.ndarray:
    """
    Replace NaN values with the mean of the non-NaN values.

    Args:
        arr: 1-D float array of shape (n,), at least one value not NaN

    Returns:
        New float array of shape (n,) with every NaN replaced by the mean
        (the input array is left unchanged)
    """
    # Your code here
    pass