Using transform with condition within a dataframe - pandas

Question

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?

Answer 1

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

Answer 2

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

Answer 3

I'd like to think that performance wise, it is better to build the temporary columns, rather than iterating through groupby:
grp = df.groupby(['animal', 'type'])['A']
(df
.assign(
mi = grp.transform('min'),
ma = grp.transform('max'),
new_col = lambda df: np.where(df['type'] == 'X', df['ma'], df['mi']))
.drop(columns=['mi','ma'])
)
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