Replace NaN in Pandas based on a multi-column condition

Viewed 1246

I hope you are doing well.

I need help to perform a complex "NaN replace" on my dataframe.

What is the best way to replace NaN values in a pandas column, based on a mode of other column values filtered by other columns?

Let me illustrate my problem:

import random
import numpy as np
import pandas as pd
data = {'Region': [1,1,1,2,2,2,1,2,2,2,2,1,1,1,2,1], 'Country': ['a','a', 'a', 'a', 'a','a', 'a', 'a', 'b', 'b', 'b', 'b','b','b','b','b'], 'GDP' : [100,100,101,105,105,110,np.nan,np.nan,200,200,100,150,100,150,np.nan,np.nan]}
df = pd.DataFrame.from_dict(data)

df:

     Region Country GDP
0        1       a  100.0
1        1       a  100.0
2        1       a  101.0
3        2       a  105.0
4        2       a  105.0
5        2       a  110.0
6        1       a    NaN
7        2       a    NaN
8        2       b  200.0
9        2       b  200.0
10       2       b  100.0
11       1       b  150.0
12       1       b  100.0
13       1       b  150.0
14       2       b    NaN
15       1       b    NaN

I would like to replace the nan values of the GDP column with the mode of other GDP values for the same country and region.

In the case of the NaN value of the GDP column of index 6, I wish to replace it with 100 (as it is the mode for GDP values for Region 1 & Country a)

The desired output should look like this:

    Region Country  GDP
0        1       a  100
1        1       a  100
2        1       a  101
3        2       a  105
4        2       a  105
5        2       a  110
6        1       a  100
7        2       a  105
8        2       b  200
9        2       b  200
10       2       b  100
11       1       b  150
12       1       b  100
13       1       b  150
14       2       b  200
15       1       b  150 

Thank you for your help, I hope you have an excellent day!

1 Answers

Pandas' fillna allows for filling missing values from another series. So we need another series that contains the mode of each Country/Region at the corresponding indices.

To get this series, we can use Pandas' groupby().transform() operation. It groups the dataframe, and then broadcasts the results back to the original shape.

If we use this operation with mode as is, it will give an error. Mode can return multiple values, preventing pandas from broadcasting the values back to the original shape. So we need to force it to return a single value, so just pick the first one (or last one, or whichever).

df["GDP"].fillna(
    df.groupby(["Country", "Region"])["GDP"].transform(
        lambda x: x.mode()[0]
    )
)
Related