如何用Numpy或其他库高效实现Pandas分组取每组前10行？

方法1：Pandas原生高效方案（推荐） #

放弃apply，改用groupby.cumcount()生成分组内的行序号，再通过布尔索引筛选，这是Pandas内置的向量化操作，速度远快于循环式的apply：

k = 10
mask = df.groupby('id').cumcount() < k
result = df[mask]

原理是给每个ID分组内的行分配从0开始的连续序号，只需保留序号小于10的行即可，全程无Python层面的循环，性能拉满。

方法2：纯Numpy实现 #

如果需要完全基于Numpy的方案，可以按以下步骤操作：

import numpy as np
import pandas as pd

# 先按ID排序，确保同ID的行连续排列
sorted_df = df.sort_values('id')
ids = sorted_df['id'].to_numpy()
float_data = sorted_df.drop('id', axis=1).to_numpy()

# 获取每个唯一ID的起始索引和出现次数
unique_ids, start_indices, counts = np.unique(ids, return_index=True, return_counts=True)
# 计算每个ID需要保留的行数（最多10行）
keep_rows = np.minimum(counts, 10)
# 拼接所有需要保留的行索引
keep_indices = np.concatenate([np.arange(start, start + num) for start, num in zip(start_indices, keep_rows)])

# 重构结果DataFrame
result = pd.DataFrame(float_data[keep_indices], columns=sorted_df.columns[1:])
result['id'] = ids[keep_indices]

这种方法完全依赖Numpy的底层向量化计算，在数据量极大时性能与Pandas原生方案接近，但代码相对繁琐，适合必须使用纯Numpy的场景。

方法3：大数据量场景用Dask #

如果你的数据集大到无法装入内存，可以用Dask进行并行处理：

import dask.dataframe as dd

# 将Pandas DataFrame转为Dask DataFrame
ddf = dd.from_pandas(df, npartitions=4)
# 分组取前10行后转回Pandas
result = ddf.groupby('id').head(10).compute()

Dask会自动将数据分片并并行处理，适合TB级别的超大规模数据集，小数据场景反而会有额外的调度开销。

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