Find which related column the user has-CodePudding

I have a problem. I have a user who lives in a country and may have one or more interests.

I just want to know what code this user has. For example, the user who lives in Germany and has Gaming and Swimming as interests has the codes 03 and 02. I would like to output these as well.

Dataframe

   country  interest code
0  Germany     Sport   01
1  Germany  Swimming   02
2  Germany    Gaming   03
3  Hungary     Sport   11
4  Hungary  Swimming   12

Code

import pandas as pd

d = {'country': ['Germany', 'Germany', 'Germany', 'Hungary', 'Hungary'],
     'interest': ['Sport', 'Swimming', 'Gaming', 'Sport', 'Swimming'],
     'code': ['01', '02','03', '11', '12'],}
df = pd.DataFrame(data=d)

print(df)

user_country = 'Germany'
user_interest = ['Sport',]

df_country = df[df['country'] == user_bundesland]
for interest in df_country['interest'].tolist():
    if(interest in user_interesse):
      print(interest)

What I want (example)

user_country = 'Germany'
user_interest = ['Sport',]


[OUT] [01]


user_country = 'Hungary'
user_interest = ['Sport', 'Swimming']


[OUT] [11, 12]

CodePudding user response：

If user_country is a string and user_interest is a list of strings, you can subset your df by the country and list of one or more interests:

user_country = 'Germany'
user_interest = ['Sport','Swimming']

user_codes = df[(df['country'] == user_country) & (df['interest'].isin(user_interest))]['code'].tolist()

The output of user_codes is:

>>> user_codes
['01', '02']

CodePudding user response：

Using df.query:

user_codes = df.query('country == "Hungary" and 
                      (interest == "Sport" or interest == "Swimming")  \
         ')['code'].tolist()

user_codes

>> [11, 12]