Pandas 向量化：计算每组满足条件的分数

Question

假设我们有一张客户及其支出表。import pandas as pddf = pd.DataFrame({    "Name":  ["Alice", "Bob", "Bob", "Charles"],    "Spend": [3, 5, 7, 9]})LIMIT = 6对于每个客户，我们可以使用以下apply方法计算他的支出中大于 6 的部分：df.groupby("Name").apply(    lambda grp: len(grp[grp["Spend"] > LIMIT]) / len(grp))NameAlice      0.0Bob        0.5Charles    1.0然而，该apply方法只是一个循环，如果有很多客户，它会很慢。问题：有没有更快的方法，大概是使用矢量化？从 0.23.4 版本开始， SeriesGroupBy 不支持比较运算符：(df.groupby("Name") ["Spend"] > LIMIT).mean()TypeError: '>' not supported between instances of 'SeriesGroupBy' and 'int'下面的代码导致 Alice 的值为空：df[df["Spend"] > LIMIT].groupby("Name").size() / df.groupby("Name").size()NameAlice      NaNBob        0.5Charles    1.0下面的代码给出了正确的结果，但它要求我们要么修改表格，要么制作副本以避免修改原始表格。df["Dummy"] = 1 * (df["Spend"] > LIMIT)df.groupby("Name") ["Dummy"] .sum() / df.groupby("Name").size()

Answer 1

Groupby 不使用矢量化，但它具有使用 Cython 优化的聚合函数。

你可以取平均值：

(df["Spend"] > LIMIT).groupby(df["Name"]).mean()

df["Spend"].gt(LIMIT).groupby(df["Name"]).mean()

或者用div0 替换 NaN：

df[df["Spend"] > LIMIT].groupby("Name").size() \

.div(df.groupby("Name").size(), fill_value = 0)

df["Spend"].gt(LIMIT).groupby(df["Name"]).sum() \

.div(df.groupby("Name").size(), fill_value = 0)

以上每个都会产生

Name

Alice      0.0

Bob        0.5

Charles    1.0

dtype: float64

表现

取决于每个条件过滤的行数和行数，因此最好在真实数据上进行测试。

np.random.seed(123)

N = 100000

df = pd.DataFrame({

    "Name":  np.random.randint(1000, size = N),

    "Spend": np.random.randint(10, size = N)

})

LIMIT = 6

In [10]: %timeit df["Spend"].gt(LIMIT).groupby(df["Name"]).mean()

6.16 ms ± 332 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [11]: %timeit df[df["Spend"] > LIMIT].groupby("Name").size().div(df.groupby("Name").size(), fill_value = 0)

6.35 ms ± 95.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [12]: %timeit df["Spend"].gt(LIMIT).groupby(df["Name"]).sum().div(df.groupby("Name").size(), fill_value = 0)

9.66 ms ± 365 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# RafaelC comment solution

In [13]: %timeit df.groupby("Name")["Spend"].apply(lambda s: (s > LIMIT).sum() / s.size)

400 ms ± 27.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [14]: %timeit df.groupby("Name")["Spend"].apply(lambda s: (s > LIMIT).mean())

328 ms ± 6.12 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

这个 NumPy 解决方案是矢量化的，但有点复杂：

In [15]: %%timeit

    ...: i, r = pd.factorize(df["Name"])

    ...: a = pd.Series(np.bincount(i), index = r)

    ...: 

    ...: i1, r1 = pd.factorize(df["Name"].values[df["Spend"].values > LIMIT])

    ...: b = pd.Series(np.bincount(i1), index = r1)

    ...: 

    ...: df1 = b.div(a, fill_value = 0)

    ...: 

5.05 ms ± 82.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)