Function for replacing missing values with median from pivot table

Question

My goal is to write a function to replace missing values in the 'total_income' column with the median 'total_income' provided by the pivot table, using the row's 'education' and 'income_type' to index the pivot table. I want to populate using these medians so that the values are as optimal as they can be. Here is what I am testing:

#This is the first 5 rows of the dataframe as a dictionary:

{'index': {0: 0, 1: 1, 2: 2, 3: 3, 4: 4},
 'children': {0: 1, 1: 1, 2: 0, 3: 3, 4: 0},
 'days_employed': {0: 8437.673027760233,
  1: 4024.803753850451,
  2: 5623.422610230956,
  3: 4124.747206540018,
  4: 340266.07204682194},
 'dob_years': {0: 42, 1: 36, 2: 33, 3: 32, 4: 53},
 'education': {0: "bachelor's degree",
  1: 'secondary education',
  2: 'secondary education',
  3: 'secondary education',
  4: 'secondary education'},
 'education_id': {0: 0, 1: 1, 2: 1, 3: 1, 4: 1},
 'family_status': {0: 'married',
  1: 'married',
  2: 'married',
  3: 'married',
  4: 'civil partnership'},
 'family_status_id': {0: 0, 1: 0, 2: 0, 3: 0, 4: 1},
 'gender': {0: 'F', 1: 'F', 2: 'M', 3: 'M', 4: 'F'},
 'income_type': {0: 'employee',
  1: 'employee',
  2: 'employee',
  3: 'employee',
  4: 'retiree'},
 'debt': {0: 0, 1: 0, 2: 0, 3: 0, 4: 0},
 'total_income': {0: 'NaN',
  1: 17932.802,
  2: 'NaN',
  3: 'NaN',
  4: 'NaN'},
 'purpose': {0: 'purchase of the house',
  1: 'car purchase',
  2: 'purchase of the house',
  3: 'supplementary education',
  4: 'to have a wedding'},
 'age_group': {0: 'adult',
  1: 'adult',
  2: 'adult',
  3: 'adult',
  4: 'older adult'}}

#this is the first 5 row of the pivot table as a dictionary:
{"bachelor's degree": {('adult', 'business'): 33702.312999999995,
  ('adult', 'civil servant'): 28178.519,
  ('adult', 'employee'): 27096.019,
  ('adult', 'paternity / maternity leave'): nan,
  ('adult', 'retiree'): 32491.072},
 'graduate degree': {('adult', 'business'): nan,
  ('adult', 'civil servant'): 17822.756999999998,
  ('adult', 'employee'): 25161.5835,
  ('adult', 'paternity / maternity leave'): nan,
  ('adult', 'retiree'): nan},
 'primary education': {('adult', 'business'): 27305.901,
  ('adult', 'civil servant'): 30545.949,
  ('adult', 'employee'): 19887.121,
  ('adult', 'paternity / maternity leave'): nan,
  ('adult', 'retiree'): 15412.875},
 'secondary education': {('adult', 'business'): 26144.03,
  ('adult', 'civil servant'): 21740.913,
  ('adult', 'employee'): 22166.627,
  ('adult', 'paternity / maternity leave'): 8612.661,
  ('adult', 'retiree'): 21219.068},
 'some college': {('adult', 'business'): 31437.118,
  ('adult', 'civil servant'): 29281.5325,
  ('adult', 'employee'): 27035.721,
  ('adult', 'paternity / maternity leave'): nan,
  ('adult', 'retiree'): 7015.422500000001}}

def fill_income(row):
    total_income = row['total_income']
    age_group =  row['age_group']
    income_type = row['income_type']
    education = row['education']
    table = df.pivot_table(index=['age_group','income_type' ], columns='education', values='total_income', aggfunc='median')


    if total_income == 'NaN':
        if age_group =='adult':
            return table.loc[education, income_type]
        if age_group =='older adult':
            return table.loc[education, income_type]
        if age_group =='younger adult':
            return table.loc[education, income_type]

#an anticipated outcome
 
print(fill_income(df.loc(1))
22166.6270

Thanks for your help!

user_12833166 · Answer

Annotated code (no need to `pivot`)

# Change the dtype to numeric
df['total_income'] = df['total_income'].astype(float)

# Calculated median per unique age_group, education and income_type
median = df.groupby(['age_group', 'education', 'income_type'])['total_income'].transform('median')

# Fill the NaN value with median
df['total_income'] = df['total_income'].fillna(median)

   index  children  days_employed  dob_years            education  education_id      family_status  family_status_id gender income_type  debt  total_income                  purpose    age_group
0      0         1    8437.673028         42    bachelor's degree             0            married                 0      F    employee     0           NaN    purchase of the house        adult
1      1         1    4024.803754         36  secondary education             1            married                 0      F    employee     0     17932.802             car purchase        adult
2      2         0    5623.422610         33  secondary education             1            married                 0      M    employee     0     17932.802    purchase of the house        adult
3      3         3    4124.747207         32  secondary education             1            married                 0      M    employee     0     17932.802  supplementary education        adult
4      4         0  340266.072047         53  secondary education             1  civil partnership                 1      F     retiree     0           NaN        to have a wedding  older adult

Function for replacing missing values with median from pivot table

1 Answers

Annotated code (no need to pivot)

Annotated code (no need to `pivot`)