如何使用多个条件在 Pandas 中添加新列

Question

I have a dataframe that looks like this:

destination_zip destination_state
502111387        IA
388588179        MS           
T2A2L9           AB                  
891              AUK     
774653028        TX   

I am trying to write a code that will be adding a new column as destination_country to my dataframe, something like this:

destination_zip destination_state  destination_country
502111387        IA                 US
388588179        MS                 US 
T2A2L9           AB                 CA
891              AUK                NZ
774653028        TX                 US

what I have tried so far is:

df.loc[df['destination_state']=='TX', df['destination_country']]= 'US'
df.loc[df['destination_state']=='IA', df['destination_country']]= 'US'
df.loc[df['destination_state']=='MS', df['destination_country']]= 'US'
df.loc[df['destination_state']=='AUK', df['destination_country']]= 'NZ'
df.loc[df['destination_state']=='AB', df['destination_country']]= 'CA'

but this is not way too long to work with, I wanted something that would be based on multiple conditions in a single line of code, something like this:

df.loc[df['destination_state']=='TX','IA','MS' , df['destination_country']]= 'US'

but this code is not working, can anyone help me with this? My dataframe has 7k rows, that's why I wanted something with multiple conditons. I am using juypter notebook, python-3

Answer 1

Here is how I would proceed:

# this format is easy to maintain
countries = {'US': ['IA', 'MS', 'TX'],
             'CA': ['AB'],
             'NZ': ['AUK'],
             }

# transform it to the inverse
states = {v:k for k,vals in countries.items() for v in vals}

# map states -> country
df['destination_country'] = df['destination_state'].map(states)

output:

  destination_zip destination_state destination_country
0       502111387                IA                  US
1       388588179                MS                  US
2          T2A2L9                AB                  CA
3             891               AUK                  NZ
4       774653028                TX                  US