Pandas Tricks for Imputing Missing Data | by Sadrach Pierre, Ph.D. | Towards Data Science

One of the biggest challenges data scientists face is dealing with missing data. In this post, we will discuss how to impute missing numerical and categorical values using Pandas.
Let’s get started!
For our purposes, we will be working with the Wine Magazine Dataset, which can be found here.
To start, let’s read the data into a Pandas data frame:
import pandas as pd df = pd.read_csv("winemag-data-130k-v2.csv")
Next, let’s print the first five rows of data using the ‘.head()’ method:
print(df.head())
Since we are interested in imputing missing values, it would be useful to see the distribution in missing values across columns.
We can display missing value information with the ‘.info()’ method. This displays the number of non-null values in each column:
print(df.info())
We can also use the ‘.isnull()’ and ‘.sum()’ methods to calculate the number of missing values in each column:
print(df.isnull().sum())
We see that the resulting Pandas series shows the missing values for each of the columns in our data. The ‘price’ column contains 8996 missing values. We can replace these missing values using the ‘.fillna()’ method. For example, let’s fill in the missing values with the mean price:
df['price'].fillna(df['price'].mean(), inplace = True)print(df.isnull().sum())
We see that the ‘price’ column no longer has missing values. Now, suppose we wanted to make a more accurate imputation. A good guess would be to replace missing values in the price column with the mean prices within the countries the missing values belong. For example, if we consider missing wine prices for Italian wine, we can replace these missing values with the mean price of Italian wine. To proceed, let’s look at the distribution in ‘country’ values. We can use the ‘Counter’ method from the collections module to do so:
from collections import Counterprint(Counter(df['country']))
Let’s take a look at wine made in the ‘US’. We can define a data frame containing only ‘US’ wines:
df_US = df[df['country']=='US']
Now, let’s print the number of missing values:
print(df_US.isnull().sum())
We see that there are 239 missing ‘price’ values in the ‘US’ wines data. To fill in the missing values with the mean corresponding to the prices in the US we do the following:
df_US['price'].fillna(df_US['price'].mean(), inplace = True)
Now suppose we wanted to do this for the missing price values in each country. First, let’s impute missing country values:
df['country'].fillna(df['country'].mode()[0], inplace = True)
Next, within a for-loop we can define country-specific data frames:
for i in list(set(df['country'])):    df_country = df[df['country']== country]
Next, we can fill the missing values in these country-specific data frames with their respective mean prices:
for i in list(set(df['country'])):    df_country = df[df['country']== country]    df_country['price'].fillna(df_country['price'].mean(),inplace = True)
We then append the result to a list we’ll call “frames”
frames = []for i in list(set(df['country'])):    df_country = df[df['country']== country]    df_country['price'].fillna(df_country['price'].mean(),inplace = True)    frames.append(df_country)
Finally, we concatenate the resulting list of data frames:
frames = []for i in list(set(df['country'])):    df_country = df[df['country']== i]    df_country['price'].fillna(df_country['price'].mean(),inplace = True)    frames.append(df_country)    final_df = pd.concat(frames)
Now, if we print the number of missing price values before imputation we get:
print(df.isnull().sum())    
And after imputation:
print(final_df.isnull().sum())
We see all but one of the missing values have been imputed. This corresponds to wines in Egypt which has no price data. We can fix this by checking the length of the data frame within the for loop and only imputing with the country-specific mean if the length is greater than one. If the length is equal to 1 we impute with the mean across all countries:
frames = []for i in list(set(df['country'])):    df_country = df[df['country']== i]    if len(df_country) > 1:            df_country['price'].fillna(df_country['price'].mean(),inplace = True)            else:        df_country['price'].fillna(df['price'].mean(),inplace = True)    frames.append(df_country)        final_df = pd.concat(frames)
Printing the result we see that all of the values have been imputed for the ‘price’ column:
print(final_df.isnull().sum())
We can define a function that generalizes this logic. Our function will take variables corresponding to a numerical column and categorical column:
def impute_numerical(categorical_column, numerical_column):    frames = []    for i in list(set(df[categorical_column])):        df_category = df[df[categorical_column]== i]        if len(df_category) > 1:                df_category[numerical_column].fillna(df_category[numerical_column].mean(),inplace = True)                else:            df_category[numerical_column].fillna(df[numerical_column].mean(),inplace = True)        frames.append(df_category)            final_df = pd.concat(frames)    return final_df
We perform imputation using our function by executing the following:
impute_price  = impute_numerical('country', 'price')print(impute_price.isnull().sum())
Let’s also verify that the shapes of the original and imputed data frames match
print("Original Shape: ", df.shape)print("Imputed Shape: ", impute_price.shape)
Similarly, we can define a function that imputes categorical values. This function will take two variables corresponding columns with categorical values.
def impute_categorical(categorical_column1, categorical_column2):    cat_frames = []    for i in list(set(df[categorical_column1])):        df_category = df[df[categorical_column1]== i]        if len(df_category) > 1:                df_category[categorical_column2].fillna(df_category[categorical_column2].mode()[0],inplace = True)                else:            df_category[categorical_column2].fillna(df[categorical_column2].mode()[0],inplace = True)        cat_frames.append(df_category)            cat_df = pd.concat(cat_frames)    return cat_df
We can impute missing ‘taster_name’ values with the mode in each respective country:
impute_taster = impute_categorical('country', 'taster_name')print(impute_taster.isnull().sum())
We see that the ‘taster_name’ column now has zero missing values. Again, let’s verify that the shape matches with the original data frame:
print("Original Shape: ", df.shape)print("Imputed Shape: ", impute_taster.shape)
I’ll stop here but feel free to play around with the data and code yourself.
To summarize, in this post we discussed how to handle missing values using the Pandas library. First, we discussed how to impute missing numerical values with the mean value across the data. We then looked at how to make category-specific numerical imputations. Finally, we showed how to impute missing categorical values with the mode corresponding to another categorical column. I hope you found this post useful/interesting. The code in this post is available on GitHub. Thank you for reading!