加速 pandas groupby 聚合

Question

我有一个如下所示的数据集：

id     count
A.     2
A.     10
B.     3
B.     13
B.     11
C.     2
C.     3

我需要为每个 ID 汇总 'count' 个值的总和和平均值。使用 pandas，我正在做这样的事情：

stats = df.groupby('id').agg({'agg': ['sum', 'mean']})

我尝试使用 numpy 数组对此进行优化：

counts = df['count'].values
ids = df['id'].values
all_ids = df['id'].unique()
stats = [(i, np.mean(df[ids==i]), np.sum(df[ids==i])) for i in all_ids]

然而，这最终比相应的 pandas 实施花费更多时间。有没有办法加快速度（也许使用一些 numpy 功能？）

Answer 1

仅当您的数据框按 id 排序时：

# a = df.sort_values('id').values
a = df.values
a = a[a[:, 0].argsort()]
groups, indexes = np.unique(a[:, 0], return_index=True)
values = np.split(a[:, 1], indexes[1:])
r = np.array([(np.mean(v), np.sum(v)) for v in values])

我想你可以用np.view来排序，但我不知道怎么用。

>>> groups
array(['A.', 'B.', 'C.'], dtype=object)

>>> indexes
array([0, 2, 5])

>>> values
[array([2, 10], dtype=object),
 array([3, 13, 11], dtype=object),
 array([2, 3], dtype=object)]

>>> r
array([[ 6. , 12. ],
       [ 9. , 27. ],
       [ 2.5,  5. ]])

性能

对于 1,000,000 个值和 1000 个不同的组

id_ = np.random.randint(1, 1001, 1000000)
count = np.random.randint(1, 60, 1000000)
df = pd.DataFrame({'id': id_, 'count': count})

%timeit np.array([(np.mean(v), np.sum(v)) for v in values])
18 ms ± 465 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

加速 pandas groupby 聚合

Speeding up pandas groupby aggregation

python

numpy

aggregate

pandas

pandas-groupby