Advanced Imputation

import pandas as pd
import numpy as np

np.random.seed(42)

# Patient records — three columns with different types of missingness
df = pd.DataFrame({
    'patient_id':  range(1, 21),
    'age':         [34, 67, 45, 72, 28, 58, 41, 63, 37, 55,
                    29, 70, 48, 65, 31, 61, 44, 68, 52, 39],
    'bmi':         [24.1, np.nan, 28.3, np.nan, 22.5, np.nan, 26.8, np.nan,
                    23.9, 27.1, 21.8, np.nan, 29.4, np.nan, 22.1, np.nan,
                    25.6, np.nan, 28.8, 23.2],
    # BMI missing for patients aged 58+ — MAR (linked to age)
    'systolic_bp': [118, 145, np.nan, 152, 112, np.nan, 128, 141, np.nan,
                    138, 115, 149, np.nan, 144, 110, np.nan, 126, 147, 132, 119],
    # systolic_bp missing randomly due to equipment failure — MCAR
    'income_band': [3, np.nan, 2, np.nan, 2, 3, 1, np.nan, 2, 3,
                    1, np.nan, 2, np.nan, 1, 3, 2, np.nan, 3, 2],
    # income_band missing because high earners (band 3+) skip form — MNAR
    'readmitted':  [0, 1, 0, 1, 0, 1, 0, 1, 0, 1,
                    0, 1, 0, 1, 0, 1, 0, 1, 0, 0]
})

print("Missing value counts:")
print(df.isnull().sum())
print(f"\nTotal missing cells: {df.isnull().sum().sum()} of {df.shape[0]*df.shape[1]}")

What just happened?

20 missing values across three columns. The clinical director already told us why each one is missing — that context is gold. Without it, we might apply the wrong fix to each column. Now let's verify the MAR hypothesis for BMI before we impute it.

from scipy import stats

# Create a binary flag: 1 = BMI is missing, 0 = BMI is present
df['bmi_missing'] = df['bmi'].isnull().astype(int)

# Compare age distribution between rows where BMI is missing vs present
age_missing = df[df['bmi_missing'] == 1]['age']
age_present = df[df['bmi_missing'] == 0]['age']

print("=== VERIFYING BMI MISSINGNESS PATTERN ===\n")
print(f"  Age when BMI is MISSING: mean={age_missing.mean():.1f}  median={age_missing.median():.0f}")
print(f"  Age when BMI is PRESENT: mean={age_present.mean():.1f}  median={age_present.median():.0f}")
print()

# t-test: is the age difference statistically significant?
t_stat, p_value = stats.ttest_ind(age_missing, age_present)
print(f"  t-test: t={t_stat:.2f}, p={p_value:.4f}")
if p_value < 0.05:
    print(f"  → CONFIRMED: BMI missingness IS significantly linked to age (p<0.05)")
    print(f"  → This is MAR — use an imputation method that uses age as context")
else:
    print(f"  → NOT confirmed: missingness may be random (MCAR)")
    print(f"  → Simple mean/median imputation is acceptable")

# Clean up helper column
df.drop(columns='bmi_missing', inplace=True)

What just happened?

pandas' .isnull().astype(int) creates a binary flag column — 1 where BMI is missing, 0 where it's present. We then split the age column by that flag and compare means. scipy's stats.ttest_ind() tests whether the age difference between the two groups is statistically significant.

The verdict is decisive: patients with missing BMI have a mean age of 64 vs 38 for those with BMI recorded — a 26-year gap with p < 0.0001. The clinical director was right. If we impute BMI with the overall mean (around 25.4), we'd be assigning a young person's BMI to elderly patients. That would quietly introduce bias into the model. We need a method that accounts for age.

from sklearn.impute import KNNImputer

# KNNImputer fills each missing value with the average of the K nearest rows
# "nearest" = most similar based on non-missing columns
# n_neighbors=5: use the 5 most similar patients to estimate each missing BMI
knn = KNNImputer(n_neighbors=5)

# We impute on the numeric columns only — KNN requires all numbers
# We use age and systolic_bp as the context columns for BMI estimation
cols_for_knn = ['age', 'bmi', 'systolic_bp']
imputed_matrix = knn.fit_transform(df[cols_for_knn])
# fit_transform learns the patterns and returns a filled numpy array

df['bmi_knn'] = imputed_matrix[:, 1].round(1)   # column index 1 = bmi

# Compare: what would mean imputation have given vs KNN?
bmi_mean = df['bmi'].mean()
df['bmi_mean_imputed'] = df['bmi'].fillna(round(bmi_mean, 1))

print("=== KNN vs MEAN IMPUTATION — BMI ===\n")
print(f"Overall BMI mean (used by simple imputation): {bmi_mean:.1f}\n")
print(f"{'PatientID':>10}  {'Age':>5}  {'Actual BMI':>11}  {'KNN Imputed':>12}  {'Mean Imputed':>13}  {'Missing?':>9}")
print("─" * 70)

for _, row in df.iterrows():
    was_missing = "← was NaN" if pd.isna(row['bmi']) else ""
    actual = f"{row['bmi']:.1f}" if pd.notna(row['bmi']) else "NaN"
    print(f"  {int(row['patient_id']):>8}  {int(row['age']):>5}  {actual:>11}  "
          f"{row['bmi_knn']:>12.1f}  {row['bmi_mean_imputed']:>13.1f}  {was_missing}")

What just happened?

sklearn's KNNImputer takes a matrix, finds the 5 most similar complete rows for each missing value, and fills in the average. fit_transform() trains and fills in one step, returning a numpy array. We use column index 1 to extract the BMI column.

The difference is stark. Mean imputation gives every missing patient 25.4 — a young person's BMI. KNN gives patient 2 (age 67) a BMI of 28.1, patient 4 (age 72) a BMI of 28.4 — both drawn from similar elderly patients in the dataset. These are plausible, contextually appropriate estimates. The mean imputation estimates are not wrong in a way you can see immediately — but they systematically bias the model against elderly patients, which is exactly the kind of hidden error that ruins model performance without triggering an obvious error message.

from sklearn.impute import SimpleImputer

# SimpleImputer fills missing values with a single statistic
# strategy='median' is more robust than 'mean' for skewed distributions
# Median is not affected by extreme blood pressure readings

bp_imputer = SimpleImputer(strategy='median')

# .fit_transform() computes the median from non-missing rows, then fills
bp_filled = bp_imputer.fit_transform(df[['systolic_bp']])
df['systolic_bp_imputed'] = bp_filled.round(0)

bp_median = df['systolic_bp'].median()

print("=== MEDIAN IMPUTATION — SYSTOLIC BP ===\n")
print(f"Median computed from non-missing rows: {bp_median:.0f} mmHg\n")

# Show before and after for missing rows only
missing_bp = df[df['systolic_bp'].isnull()][['patient_id','age',
                                               'systolic_bp','systolic_bp_imputed']]
print(missing_bp.to_string(index=False))
print()

# Verify the imputed column has no more missing values
print(f"Missing systolic_bp after imputation: {df['systolic_bp_imputed'].isnull().sum()}")

What just happened?

sklearn's SimpleImputer(strategy='median') computes the median of the non-missing values during .fit() and fills all NaNs with that value during .transform(). We use median rather than mean because blood pressure distributions can be skewed by extreme values — a patient with a hypertensive crisis of 200+ mmHg would inflate the mean but not the median.

All six missing patients get 132 mmHg — which is the median of the observed readings. For MCAR data this is correct: the missing rows are a random sample of all patients, so filling with the population median introduces no systematic bias. Simple, defensible, done.

# For MNAR: do NOT fill with median — it would introduce systematic bias
# Instead: create a missingness flag that the model can learn from
# The fact that income is missing IS itself a signal — it likely means high earner

df['income_missing_flag'] = df['income_band'].isnull().astype(int)
# income_missing_flag = 1 means "this patient declined to report income"
# A smart model will learn: when income_missing_flag=1, readmission risk may differ

# Show the relationship between income being missing and the target
print("=== MNAR ANALYSIS — INCOME BAND ===\n")
print("Readmission rate by income disclosure:\n")

disclosed = df[df['income_missing_flag'] == 0]
withheld  = df[df['income_missing_flag'] == 1]

print(f"  Income DISCLOSED:  n={len(disclosed)}  readmission rate = "
      f"{disclosed['readmitted'].mean()*100:.0f}%")
print(f"  Income WITHHELD:   n={len(withheld)}   readmission rate = "
      f"{withheld['readmitted'].mean()*100:.0f}%")
print()

# Verify that the income_missing_flag column IS correlated with readmission
from scipy import stats
r, p = stats.pointbiserialr(df['income_missing_flag'], df['readmitted'])
print(f"  Correlation of income_missing_flag with readmitted: r={r:.3f}  p={p:.4f}")
if p < 0.05:
    print(f"  → The flag IS a significant predictor. Keep it as a feature.")
else:
    print(f"  → Weak signal, but still worth including as a low-cost feature.")

print()
print("Final column plan:")
print("  income_band:         → keep as-is (NaNs will be imputed or handled by model)")
print("  income_missing_flag: → add as a new feature column")

What just happened?

pandas' .isnull().astype(int) creates the missingness flag. stats.pointbiserialr() is the correct correlation test for a binary flag vs a continuous or binary target — equivalent to Pearson for this case.

The analysis confirms the MNAR hypothesis: patients who withheld income have a 67% readmission rate vs 29% for those who disclosed. The missingness itself is a clinical signal — these patients are different in ways that affect outcomes. Filling their income with the median (band 2) would have erased that signal. By keeping the flag, we let the model see that "income withheld" is a meaningful category, not noise.

Teacher's Note

Imputation is not cleaning. It is estimation. Every imputed value is a guess. The goal is to make a guess that is consistent with the data's actual structure — which means understanding why the value is missing before deciding how to replace it.

The most important piece of this lesson isn't the code — it's the conversation you should have with the data owner or domain expert before touching a single NaN. "Why is this value missing?" is the most important question in imputation. The code only matters once you have the answer.

Up Next · Lesson 38

Feature Engineering via EDA

Use EDA insights to build features that models can't discover on their own — interaction terms, polynomial features, and target encodings grounded in what the data actually shows you.