如何加速Pandas GroupBy聚合操作？Numpy优化方案效率低下的技术问询

我有如下结构的数据集：

我需要针对每个ID，聚合计算count字段的总和与均值。用Pandas时，我用下面的代码实现：

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

为了优化性能，我尝试改用Numpy数组实现聚合逻辑：

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

但实际运行后发现，这个Numpy实现的耗时反而超过了对应的Pandas原生实现。请问是否存在可行的加速方案（或许可借助特定Numpy功能实现）？

你的Numpy实现性能不佳的核心原因是：循环遍历每个唯一ID时，每次都执行ids==i的布尔索引，这会重复扫描整个数组，时间复杂度是O(N*K)（N是总行数，K是唯一ID的数量）。而Pandas的groupby底层是用C优化过的，这种朴素循环自然比不过它。

要让Numpy实现追上甚至超过Pandas的性能，关键是用向量化操作替代Python循环，避免重复扫描数组。这里可以利用np.unique的return_inverse参数和np.bincount来实现：

优化后的Numpy代码 #

import numpy as np
import pandas as pd

# 从DataFrame提取Numpy数组
counts = df['count'].values
ids = df['id'].values

# 获取唯一ID，以及每个元素对应的唯一ID的索引（把字符串ID映射成整数索引）
unique_ids, idx = np.unique(ids, return_inverse=True)

# 用bincount计算每个分组的总和（向量化操作，一次遍历完成）
count_sum = np.bincount(idx, weights=counts)

# 计算每个分组的元素数量，再求均值
group_sizes = np.bincount(idx)
count_mean = count_sum / group_sizes

# 组合成你需要的结果格式
stats = list(zip(unique_ids, count_sum, count_mean))

# 如果需要转换成Pandas DataFrame格式（和原生groupby结果对齐）
stats_df = pd.DataFrame({
    'sum': count_sum,
    'mean': count_mean
}, index=unique_ids)

为什么这个方法更快？ #

1. 时间复杂度降到O(N)：np.unique和np.bincount都是单次遍历数组完成计算，没有Python循环的额外开销，也不会重复扫描数据。
2. 利用Numpy的底层优化：bincount是用C实现的向量化操作，效率和Pandas的groupby底层逻辑接近，甚至在纯数值计算场景下会更高效。

性能对比 #

• 当数据集较小时，Pandas和优化后的Numpy方法差距不大；
• 当数据集很大（比如百万级以上）且唯一ID数量较多时，这个Numpy实现的性能会明显超过你的朴素循环版本，甚至能和Pandas的groupby打平或略胜一筹。

内容的提问来源于stack exchange，提问作者Karan