pandas - 在数据框中使用带有条件的变换

Question

有以下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))))

如果动物的类型是 X，我希望获得 A 列的 max ...否则在单独的列中获得 A 列的 min .

应用 lambda 和 group by 显示具有以下代码的多索引数组:

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

有没有办法将其转换为系列，可以将其作为列添加到 df...例如使用 transform？

Answer 1

这是你想要的吗？

>>> 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