median imputer function

median_imputer.py

def ImputeWithMedian(train_df, test_df, cols=["price", "year"]):
  # function to impute price and year
  # columns with the median value of each 
  # median imputer is fitted on train dataset
  # transformation done on the train and test set

  train_df = pd.DataFrame(train_df[cols])
  test_df = pd.DataFrame(test_df[cols]) 
    
  median_imputer = SimpleImputer(strategy="median")

  median_imputer.fit(train_df)
    
  imputed_train_set = median_imputer.transform(train_df)
  imputed_train_df = pd.DataFrame(imputed_train_set, columns=train_df.columns)

  imputed_test_set = median_imputer.transform(test_df)
  imputed_test_df = pd.DataFrame(imputed_test_set, columns=test_df.columns)

  return imputed_train_df, imputed_test_df