应用 Python Pandas 进行分组的更快方法？

Question

如何让 Groupby Apply 运行得更快，或者如何以不同的方式编写它？

import numpy as np
import pandas as pd

df = pd.DataFrame({'ID':[1,1,1,1,1,2,2,2,2,2],\
                   'value':[1,2,np.nan,3,np.nan,1,2,np.nan,4,np.nan]})

result = df.groupby("ID").apply(lambda x: len(x[x['value'].notnull()].index)\
                    if((len(x[x['value']==1].index)>=1)&\
                    (len(x[x['value']==4].index)==0)) else 0)

输出：

Index  0  
1      3  
2      0

我的程序现在运行得很慢。我可以让它更快吗？我过去在使用 groupby() 之前进行了过滤，但在这种情况下我没有看到一个简单的方法来做到这一点。

Answer 1

不确定这是否是您所需要的。我已经对其进行了一些分解，但是您可以轻松地对其进行方法链接以使代码更加紧凑：

df = pd.DataFrame(
    {
        "ID": [1, 1, 1, 1, 1, 2, 2, 2, 2, 2],
        "value": [1, 2, np.nan, 3, np.nan, 1, 2, np.nan, 4, np.nan],
    }
)

df["x1"] = df["value"] == 1
df["x2"] = df["value"] == 4

df2 = df.groupby("ID").agg(
    y1=pd.NamedAgg(column="x1", aggfunc="max"),
    y2=pd.NamedAgg(column="x2", aggfunc="max"),
    cnt=pd.NamedAgg(column="value", aggfunc="count"),
)

df3 = df2.assign(z=lambda x: (x['y1'] & ~x['y2'])*x['cnt'])

result = df3.drop(columns=['y1', 'y2', 'cnt'])
print(result)

这会产生

    z
ID   
1   3
2   0

Answer 2

尝试swifter：一个以最快的可用方式将任何功能有效地应用于pandas数据框或系列的包。

%%time
result = df.groupby("ID").apply(lambda x: len(x[x['value'].notnull()].index)\
                    if((len(x[x['value']==1].index)>=1)&\
                    (len(x[x['value']==4].index)==0)) else 0)

CPU times: user 7.04 ms, sys: 0 ns, total: 7.04 ms
Wall time: 6.93 ms

import swifter

%%time
result = df.swifter.groupby("ID").apply(lambda x: len(x[x['value'].notnull()].index)\
                    if((len(x[x['value']==1].index)>=1)&\
                    (len(x[x['value']==4].index)==0)) else 0)

CPU times: user 5.93 ms, sys: 114 µs, total: 6.04 ms
Wall time: 6 ms