How to efficiently assign a column value based on the occurance of string in Pandas?

Viewed 48

The objective is to assign the value in the column GP based on a look up table (i.e., group_post_dict). If the text in GP contain any of the string in the lookup, then the column GP will be assign based proper lookup and group_ pair.

The following code demonstrate better,

group_post_dict = {1: {'lookup': ['pet dog', 'pet cat'], 'group_': 'pet'},
                   2: {'lookup': ['flow red', 'flow yellow'], 'group_': 'flow'},
                   3: {'lookup': ['b year', 'c year'], 'group_': 'year'}}

data_list=['pet dog', 'pet cat','flow red', 'flow yellow','b year', 'c year']
df = DataFrame ( data_list)
df.columns =['PU']


for key, value in group_post_dict.items ():
    mask_sv = df ['PU'].str.contains (
        r'\b(?:{})\b'.format ( '|'.join ( value ['lookup'] ) ) )
    df.loc [mask_sv, 'GP'] = value ['group_']

How, the number of loop the for loop need to iterate increase with the number of lookup and group_ pair.

May I know whether this code can be improved further?

p.s., sorry for the lake of better title

2 Answers

Use Series.map with flatten values of dictionary for flatten dict:

d= {x:v['group_'] for k, v in group_post_dict.items() for x in v['lookup']}
print (d)
{'pet dog': 'pet', 'pet cat': 'pet', 'flow red': 'flow', 
 'flow yellow': 'flow', 'b year': 'year', 'c year': 'year'}

df['GP'] = df['PU'].map(d)
 
print (df)
            PU    GP
0      pet dog   pet
1      pet cat   pet
2     flow red  flow
3  flow yellow  flow
4       b year  year
5       c year  year

Construct dataframe from group_post_dict, explode and map

df['GP'] = df.PU.map(pd.DataFrame(group_post_dict).T
                       .explode('lookup').set_index('lookup')['group_'])

Out[158]:
            PU    GP
0      pet dog   pet
1      pet cat   pet
2     flow red  flow
3  flow yellow  flow
4       b year  year
5       c year  year
Related