INNER CODE UNIT · Python
impute_NA_with_avg
ashishpatel26/Amazing-Feature-Engineering · feature_cleaning/missing_data.py:71
def impute_NA_with_avg(data,strategy='mean',NA_col=[]):
"""
replacing the NA with mean/median/most frequent values of that variable.
Note it should only be performed over training set and then propagated to test set.
"""
data_copy = data.copy(deep=True)
for i in NA_col:
if data_copy[i].isnull().sum()>0:
if strategy=='mean':
data_copy[i+'_impute_mean'] = data_copy[i].fillna(data[i].mean())
elif strategy=='median':
data_copy[i+'_impute_median'] = data_copy[i].fillna(data[i].median())
elif strategy=='mode':
data_copy[i+'_impute_mode'] = data_copy[i].fillna(data[i].mode()[0])
else:
warn("Column %s has no missing" % i)
return data_copy