Skip to main content
 首页 » 编程设计

python之使用 pandas 按列分组,然后根据条件创建新列

2025年05月04日12haluo1

我需要用 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 

这是一个示例,说明性的 pandas 数据框:

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”)。

请您参考如下方法:

对条件逻辑表达式中的真值求和:

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,您也可以在 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