Broadcast value to dataframe group by condition

Viewed 514

Dataframe:

STUD_ID   CLASS   GRADE  CATEGORY 
1         'Sci'    max  'Alpha' 
2         'Sci'    7    'Omega' 
3         'Sci'    9    'Alpha' 
4         'Sci'    3    'Alpha' 
5         'Sci'    9    'Alpha' 
7         'eng'    max  'Gamma' 
8         'eng'    5    'Gamma' 
9         'eng'    3    'Epsilon'
10        'Art'    3    'Lambda' 
11        'Art'    8    'Lambda' 
12        'Art'    max  undefined 
13        'Art'    1    undefined 
14        'Art'    5    undefined 

I want to create a new column CATEGORY_1 that broadcasts the CATEGORY value (where the GRADE is max) to their respective CLASS group. See result df below.

STUD_ID   CLASS   GRADE  CATEGORY CATEGORY_1
1         'Sci'    max  'Alpha' .   'Alpha'
2         'Sci'    7    'Omega' .   'Alpha'
3         'Sci'    9    'Alpha' .   'Alpha'
4         'Sci'    3    'Alpha' .   'Alpha'
5         'Sci'    9    'Alpha' .   'Alpha'
7         'eng'    max  'Gamma' .   'Gamma'
8         'eng'    5    'Gamma' .   'Gamma'
9         'eng'    3    'Epsilon'   'Gamma'
10        'Art'    3    'Lambda' .  undefined
11        'Art'    8    'Lambda' .  undefined
12        'Art'    max  undefined . undefined
13        'Art'    1    undefined . undefined
14        'Art'    5    undefined . undefined

I've tried using groupby + transform but I cannot figure out how to broadcast only the CATEGORY value where GRADE is max.

4 Answers

Because you mention groupby + transform:

It's difficult because your transformation depends on more than one column. This requires your transformation function to accept both the Series, and the entire DataFrame, which you can then subset to within group (based on the Series index). You should have a unique index for this to work properly.

def get_max_cat(grade, df):
    s = grade=='max'
    return df.loc[s[s].index, 'CATEGORY'].item()  # Assumes always one max

df['Category_1'] = df.groupby('CLASS').GRADE.transform(get_max_cat, df=df)

If I understand correctly, you can achieve this with a dictionary mapping:

df['CATEGORY_1'] = df.CLASS.map(pd.Series(df[df.GRADE == 'max']['CATEGORY'].values,index=df[df.GRADE == 'max']['CLASS']).to_dict())

Output:

         CLASS GRADE   CATEGORY CATEGORY_1
STUD_ID                                   
1        'Sci'   max    'Alpha'    'Alpha'
2        'Sci'     7    'Omega'    'Alpha'
3        'Sci'     9    'Alpha'    'Alpha'
4        'Sci'     3    'Alpha'    'Alpha'
5        'Sci'     9    'Alpha'    'Alpha'
7        'eng'   max    'Gamma'    'Gamma'
8        'eng'     5    'Gamma'    'Gamma'
9        'eng'     3  'Epsilon'    'Gamma'
10       'Art'     3   'Lambda'  undefined
11       'Art'     8   'Lambda'  undefined
12       'Art'   max  undefined  undefined
13       'Art'     1  undefined  undefined
14       'Art'     5  undefined  undefined

you can first select the rows with the value 'max' in the column GRADE and then merge on CLASS such as:

df = df.merge( df[df.GRADE == 'max', ['CLASS','CATEGORY']], 
               on='CLASS, how='left',suffixes=('','_1'))

Another way is using where to turn CATEGORY of GRADE not max into NaN. Then, apply ffill and bfill to each group of CLASS

df['CATEGORY_1'] = df.CATEGORY.where(df.GRADE.eq('max')) \
                     .groupby(df.CLASS).apply(lambda x: x.ffill().bfill())


Out[1493]:
    STUD_ID  CLASS GRADE   CATEGORY CATEGORY_1
0         1  'Sci'   max    'Alpha'    'Alpha'
1         2  'Sci'     7    'Omega'    'Alpha'
2         3  'Sci'     9    'Alpha'    'Alpha'
3         4  'Sci'     3    'Alpha'    'Alpha'
4         5  'Sci'     9    'Alpha'    'Alpha'
5         7  'eng'   max    'Gamma'    'Gamma'
6         8  'eng'     5    'Gamma'    'Gamma'
7         9  'eng'     3  'Epsilon'    'Gamma'
8        10  'Art'     3   'Lambda'  undefined
9        11  'Art'     8   'Lambda'  undefined
10       12  'Art'   max  undefined  undefined
11       13  'Art'     1  undefined  undefined
12       14  'Art'     5  undefined  undefined
Related