Python pandas, how to replace dataframe NA value by the value looked up in an another dataframe?

Viewed 37

Suppose that I have a dataframe:

>>> import pandas as pd
>>> import numpy as np
>>> rand = np.random.RandomState(42)
>>> data_points = 10
>>> dates = pd.date_range('2020-01-01', periods=data_points, freq='D')
>>> state_city = [('USA', 'Washington'), ('France', 'Paris'), ('Germany', 'Berlin')]
>>>
>>> df = pd.DataFrame()
>>> for _ in range(data_points):
...     state, city = state_city[rand.choice(len(state_city))]
...     df_row = pd.DataFrame(
...         {
...             'time' : rand.choice(dates),
...             'state': state,
...             'city': city,
...             'val1': rand.randint(0, data_points),
...             'val2': rand.randint(0, data_points)
...         }, index=[0]
...     )
...
...     df = pd.concat([df, df_row], ignore_index=True)
...
>>> df = df.sort_values(['time', 'state', 'city']).reset_index(drop=True)
>>> df.loc[rand.randint(0, data_points, size=rand.randint(1, 3)), ['state']] = pd.NA
>>> df.loc[rand.randint(0, data_points, size=rand.randint(1, 3)), ['city']] = pd.NA
>>> df.val1 = df.val1.where(df.val1 < 5, pd.NA)
>>> df.val2 = df.val2.where(df.val2 < 5, pd.NA)
>>>
>>> df
        time    state        city  val1  val2
0 2020-01-03      USA  Washington     4     2
1 2020-01-04   France        <NA>  <NA>     1
2 2020-01-04  Germany      Berlin  <NA>     4
3 2020-01-05  Germany      Berlin  <NA>  <NA>
4 2020-01-06   France       Paris     1     4
5 2020-01-06  Germany      Berlin     4     1
6 2020-01-08  Germany      Berlin     4     3
7 2020-01-10  Germany      Berlin     2  <NA>
8 2020-01-10     <NA>  Washington  <NA>  <NA>
9 2020-01-10     <NA>  Washington     2  <NA>
>>>

As you can see there are some values. I would like to impute values for state/city as much as I could. In order to do so I will generate dataframe that can help.

>>> known_state_city = df[['state', 'city']].dropna().drop_duplicates()
>>> known_state_city
     state        city
0      USA  Washington
2  Germany      Berlin
4   France       Paris

OK, now we have all state/city combinations.

How to use known_state_city dataframe in order to fill empty state when city is known? I can find empty states that have populated city with:

>>> df.loc[df.state.isna() & df.city.notna(), 'city']
8    Washington
9    Washington
Name: city, dtype: object

But how to replace Washington with state from known_state_city not breaking index values (8 and 9) in order to replace df.state values? If I don't have all combinations in known_state_city, how to update state in df with what I have?

1 Answers

We can do fillna with map twice:

# fill empty state
df['state'] = df['state'].fillna(df['city'].map(known_state_city.set_index('city')['state']))

# fill empty city
df['city'] = df['city'].fillna(df['state'].map(known_state_city.set_index('state')['city']))

Output:

         time    state        city  val1  val2
0  2020-01-03      USA  Washington   4.0   2.0
1  2020-01-04   France       Paris   NaN   1.0
2  2020-01-04  Germany      Berlin   NaN   4.0
3  2020-01-05  Germany      Berlin   NaN   NaN
4  2020-01-06   France       Paris   1.0   4.0
5  2020-01-06  Germany      Berlin   4.0   1.0
6  2020-01-08  Germany      Berlin   4.0   3.0
7  2020-01-10  Germany      Berlin   2.0   NaN
8  2020-01-10      USA  Washington   NaN   NaN
9  2020-01-10      USA  Washington   2.0   NaN
Related