使用pandas按列分组，然后根据条件新建一列

Question

我需要用 pandas 轻松重现 SQL 所做的事情：

select
    del_month
    , sum(case when off0_on1 = 1 then 1 else 0 end) as on1
    , sum(case when off0_on1 = 0 then 1 else 0 end) as off0
from a1
group by del_month
order by del_month

这是一个示例，说明性的熊猫数据框：

a1 = pd.DataFrame({'del_month':[1,1,1,1,2,2,2,2], 'off0_on1':[0,0,1,1,0,1,1,1]})

这是我用 pandas 重现上述 SQL 的尝试。第一行有效。第二行报错：

a1['on1'] = a1.groupby('del_month')['off0_on1'].transform(sum)
a1['off0'] = a1.groupby('del_month')['off0_on1'].transform(sum(lambda x: 1 if x == 0 else 0))

这是第二行的错误：

TypeError: 'function' object is not iterable

这个previous question of mine 的 lambda 函数有问题，已解决。更大的问题是如何在分组数据上重现 SQL 的“sum(case when)”逻辑。我正在寻找一个通用的解决方案，因为我需要经常做这种事情。我上一个问题中的答案建议在 lambda 函数中使用 map()，但是“off0”列的以下结果不是我需要的。 “on1”列是我想要的。整个组的答案应该是相同的（即“del_month”）。

Answer 1

简单地将条件逻辑表达式中的 True 相加：

import pandas as pd

a1 = pd.DataFrame({'del_month':[1,1,1,1,2,2,2,2], 
                   'off0_on1':[0,0,1,1,0,1,1,1]})

a1['on1'] = a1.groupby('del_month')['off0_on1'].transform(lambda x: sum(x==1))    
a1['off0'] = a1.groupby('del_month')['off0_on1'].transform(lambda x: sum(x==0))

print(a1)    
#    del_month  off0_on1  on1  off0
# 0          1         0    2     2
# 1          1         0    2     2
# 2          1         1    2     2
# 3          1         1    2     2
# 4          2         0    3     1
# 5          2         1    3     1
# 6          2         1    3     1
# 7          2         1    3     1

同样，如果方言支持它，你可以在 SQL 中做同样的事情，这是最应该的：

select
    del_month
    , sum(off0_on1 = 1) as on1
    , sum(off0_on1 = 0) as off0
from a1
group by del_month
order by del_month

要在 pandas 中复制上述 SQL，不要使用 transform，而是在 groupby().apply() 调用中发送多个聚合：

def aggfunc(x):
    data = {'on1': sum(x['off0_on1'] == 1),
            'off0': sum(x['off0_on1'] == 0)}

    return pd.Series(data)

g = a1.groupby('del_month').apply(aggfunc)

print(g)    
#            on1  off0
# del_month           
# 1            2     2
# 2            3     1

Answer 2

使用get_dummies 只需要一个groupby 调用，这样更简单。

v = pd.get_dummies(df.pop('off0_on1')).groupby(df.del_month).transform(sum)
df = pd.concat([df, v.rename({0: 'off0', 1: 'on1'}, axis=1)], axis=1)

df
   del_month  off0  on1
0          1     2    2
1          1     2    2
2          1     2    2
3          1     2    2
4          2     1    3
5          2     1    3
6          2     1    3
7          2     1    3

另外，对于聚合的情况，直接调用sum，不要使用apply：

(pd.get_dummies(df.pop('off0_on1'))
   .groupby(df.del_month)
   .sum()
   .rename({0: 'off0', 1: 'on1'}, axis=1))

           off0  on1
del_month           
1             2    2
2             1    3