Concepts 🚀 Project 🧮 Labs

Imputation — when to use mean vs median vs mode vs KNN vs MICE

Missing values are NOT a 'fill with mean' problem. Different imputation strategies for different situations. Know when to use each.

script.py v3.1

import pandas as pd
import numpy as np
from sklearn.experimental import enable_iterative_imputer  # noqa
from sklearn.impute import SimpleImputer, KNNImputer, IterativeImputer
from sklearn.ensemble import RandomForestRegressor

# ── Create a dataset with KNOWN missing patterns ──
np.random.seed(42)
n = 500
df = pd.DataFrame({
    'age': np.random.normal(40, 12, n).clip(18, 80),        # numeric, normal-ish
    'income': np.random.lognormal(10.5, 0.6, n),             # numeric, SKEWED (long tail)
    'city': np.random.choice(['NYC', 'LA', 'SF', 'Chicago'], n),  # categorical
    'credit_score': np.random.normal(680, 80, n).clip(300, 850),  # numeric, bounded
})

# Inject missing values (different patterns)
df.loc[df.sample(50, random_state=42).index, 'age'] = np.nan          # 10% MCAR
df.loc[df[df['income'] > 50000].sample(40, random_state=42).index, 'income'] = np.nan  # MNAR (rich hide income)
df.loc[df.sample(30, random_state=42).index, 'city'] = np.nan         # 6% categorical
df.loc[df.sample(80, random_state=42).index, 'credit_score'] = np.nan  # 16%

print('Missing values per column:')
print(df.isnull().sum())
print(f'\nTotal rows: {len(df)}')

# ════════════════════════════════════════════════════════════════════
#  5 IMPUTATION STRATEGIES — pick based on the DATA, not habit
# ════════════════════════════════════════════════════════════════════

print('\n' + '='*60)
print('STRATEGY 1: Mean imputation (for NORMAL distributions)')
print('='*60)
# GOOD for: age (normal-ish). BAD for: income (skewed — mean pulled by outliers).
mean_imp = SimpleImputer(strategy='mean')
df_mean = df.copy()
df_mean['age'] = mean_imp.fit_transform(df[['age']])
print(f'  age mean imputed with: {mean_imp.statistics_[0]:.2f}')
print(f'  WHY: age is roughly normal, so mean ≈ median ≈ typical value')

print('\n' + '='*60)
print('STRATEGY 2: Median imputation (for SKEWED distributions)')
print('='*60)
# GOOD for: income (skewed). BAD for: categorical.
med_imp = SimpleImputer(strategy='median')
df_med = df.copy()
df_med['income'] = med_imp.fit_transform(df[['income']])
print(f'  income median imputed with: {med_imp.statistics_[0]:.2f}')
print(f'  WHY: income is right-skewed (lognormal). Mean is pulled up by billionaires.')
print(f'  Median is robust to outliers — better represents the typical person.')

print('\n' + '='*60)
print('STRATEGY 3: Mode imputation (for CATEGORICAL)')
print('='*60)
# GOOD for: city (categorical). NEVER use mean/median on categories.
mode_imp = SimpleImputer(strategy='most_frequent')
df_mode = df.copy()
df_mode['city'] = mode_imp.fit_transform(df[['city']]).ravel()
print(f'  city mode imputed with: {mode_imp.statistics_[0]}')
print(f'  WHY: categorical has no mean. Most frequent = safest default.')

print('\n' + '='*60)
print('STRATEGY 4: KNN imputation (uses CORRELATION between features)')
print('='*60)
# GOOD for: when features are correlated. age ~ credit_score ~ income.
# BAD for: large datasets (slow — O(n²)).
knn_imp = KNNImputer(n_neighbors=5)
df_knn = df.copy()
# KNN only works on numeric — encode city first or skip it
numeric_cols = ['age', 'income', 'credit_score']
df_knn[numeric_cols] = knn_imp.fit_transform(df[numeric_cols])
print(f'  KNN imputed age, income, credit_score using 5 nearest neighbors')
print(f'  WHY: a 25-year-old is more like other 25-year-olds than like the population mean.')
print(f'  KNN finds similar rows and uses THEIR values — more accurate than mean.')

print('\n' + '='*60)
print('STRATEGY 5: MICE / IterativeImputer (models each feature using others)')
print('='*60)
# GOOD for: complex missing patterns. BAD for: speed (trains a model per column).
mice_imp = IterativeImputer(
    estimator=RandomForestRegressor(n_estimators=10, random_state=42),
    max_iter=10, random_state=42
)
df_mice = df.copy()
df_mice[numeric_cols] = mice_imp.fit_transform(df[numeric_cols])
print(f'  MICE trained a RandomForest to predict each missing value from the others')
print(f'  WHY: most sophisticated. Models the RELATIONSHIP between features.')
print(f'  E.g., missing income predicted from age + credit_score + city.')

# ════════════════════════════════════════════════════════════════════
#  COMPARISON — how different are the results?
# ════════════════════════════════════════════════════════════════════
print('\n' + '='*60)
print('COMPARISON — income imputed values (first 10 nulls):')
print('='*60)
income_nulls = df['income'].isnull()
print(f'  Median:  {df_med.loc[income_nulls, "income"].head().values}')
print(f'  KNN:     {df_knn.loc[income_nulls, "income"].head().values}')
print(f'  MICE:    {df_mice.loc[income_nulls, "income"].head().values}')

Python output click Run to execute

Click "Run" to execute the code on the left.
The output will appear here.

🧠 How it works

5 imputation strategies — pick based on the DATA: MEAN — for normal distributions (age, height). Pulled by outliers. MEDIAN — for skewed distributions (income, prices). Robust to outliers. MODE (most_frequent) — for categorical (city, color). The only option for categories. KNN — uses correlation. A 25-year-old's missing income = average of 5 similar 25-year-olds. More accurate but O(n²) — slow on large data. MICE (IterativeImputer) — trains a model per column to predict missing values from others. Most sophisticated. Best for complex patterns. Slowest. THREE missingness types you should know: MCAR (Missing Completely At Random) — safe to impute. MAR (Missing At Random) — depends on other features (KNN/MICE help). MNAR (Missing Not At Random) — depends on the missing value itself (rich hide income). Hardest.

✏️ Try these changes

Change income strategy from 'median' to 'mean'
→ mean is pulled up by billionaires — bad for skewed data
Change KNN n_neighbors from 5 to 20
→ more neighbors = smoother but less specific imputation
Add a constant imputer: SimpleImputer(strategy='constant', fill_value=-999)
→ marks missing as a separate value (good for tree models)
Add strategy='constant' fill_value='Unknown' for city
→ treats missing as its own category

← Previous lab

Preprocessing mindset — the 7 questions to ask ANY dataset

Next lab →

Encoding — OneHot vs Label vs Target vs Frequency vs Embedding