Impute Missing Values

Fill in missing values, marked as nan, with a statistic of each column.

Real tables have holes, and most models cannot accept nan. The usual fix is to replace each missing value with the mean or the median of its column. The important part is where the statistic comes from: compute it on the training data only, then use the same numbers to fill both the training set and the test set. Computing it on the test set leaks information about the test data into the model.

• strategy="mean": use the column mean, ignoring nan.
• strategy="median": use the column median, ignoring nan.
• If a training column has no values at all, fill it with 0.0.

Function signature

def impute_missing(X_train, X_test, strategy="mean"):

Example

Input

X_train = np.array([[1.0, np.nan],
                    [3.0, 4.0],
                    [np.nan, 8.0]])
X_test  = np.array([[np.nan, np.nan]])
impute_missing(X_train, X_test)

Output

(array([[1., 6.],
        [3., 4.],
        [2., 8.]]),
 array([[2., 6.]]))

Reasoning

The training column means, ignoring nan, are (1 + 3) / 2 = 2 and (4 + 8) / 2 = 6. Each nan in either set is replaced with its column's training mean.

Constraints and notes

• Both inputs are 2D float arrays with the same number of columns.
• Return new arrays. Do not change the inputs.