Using transform with condition within a dataframe

user_7978112 0
Viewed 124

Have the following df

import numpy as np
import random

i = ['dog', 'cat', 'rabbit', 'elephant'] * 20

df = pd.DataFrame(np.random.randn(len(i), 3), index=i, \
            columns=list('ABC')).rename_axis('animal').reset_index()
            
df.insert(1, 'type', pd.Series(random.choice(['X', 'Y']) \
                for _ in range(len(df))))

I would like to have the max of column A, if the type of the animal is X ... else the min of column A, in a separate column.

Apply lambda with group by shows the multi-indexed array with the following code:

g = df.groupby(['animal', 'type'])
g.apply(lambda g: np.where (g.type == 'X', g.A.max(), g.A.min()))

Is there a way to convert this to a series, that can be added to df as a column... say by using transform?

pandas transform
2 Answers

Is this what you want?

>>> df

     animal type     A     B     C
0       cat    Y  0.96 -0.02 -0.14
1       cat    Y -0.80  0.86  1.75
2       dog    X  1.13 -0.49 -1.66
3       dog    Y  0.84 -0.07  0.15
4  elephant    X  0.13 -0.54  0.73
5  elephant    Y  0.14  1.77  0.94
6    rabbit    X -0.12 -0.39  0.05
7    rabbit    X  0.58 -1.17  0.77

>>> def max_min_A(g):
        animal, type_ = g.name 
        return np.where(type_ == 'X', g.max(), g.min())

>>> df['new_col'] = df.groupby(['animal', 'type'])['A'].transform(max_min_A)

     animal type     A     B     C  new_col
0       cat    Y  0.96 -0.02 -0.14    -0.80
1       cat    Y -0.80  0.86  1.75    -0.80
2       dog    X  1.13 -0.49 -1.66     1.13
3       dog    Y  0.84 -0.07  0.15     0.84
4  elephant    X  0.13 -0.54  0.73     0.13
5  elephant    Y  0.14  1.77  0.94     0.14
6    rabbit    X -0.12 -0.39  0.05     0.58
7    rabbit    X  0.58 -1.17  0.77     0.58
user_171206920

@HarryPlotter: Thx for the name info. It is wonderful to see that the name of the group propagates as a tuple. In case one does not want to use a function, the following will work:

df.assign(new_col=g.A.transform(lambda x: np.where(x.name[1] =='X', \
            x.max(), x.min()))) 
# x.name[1] is used to select the second element of the tuple, which is `type`
user_79781120
Related