Pandas按ids分组转tolist效率优化求助

我有一个包含30-40万行的Pandas DataFrame，结构如下：

>>> df
             speed       acc                ids  row_id
0.0       0.009766  0.004053            89806.0  8819.0
1.0       0.009766  0.001790            89886.0  8818.0
2.0       0.009766  0.000006            89886.0  8834.0
3.0       0.009766  0.003188            89886.0  8835.0
4.0       0.009766  0.007101            89806.0  8836.0
...            ...       ...                ...     ...
332659.0  8.271179  0.073181            89876.0  1816.0
332660.0  8.273315  0.297774            89876.0  1813.0
332661.0  8.282776  0.130999            89876.0  1815.0
332662.0  8.285522  0.208244            89876.0  1814.0
332663.0  8.291321  0.000395            89876.0  1822.0

需要按ids字段分组，将每组转换为嵌套列表形式，示例输出如下：

out = [[[0.009766, 0.004053, 89806.0, 8819.0], [0.009766,0.007101,89806.0,8836.0]] , [[0.009766,0.001790,89886.0,8818.0], [0.009766,0.000006,89886.0,8834.0], [0.009766,0.003188,89886.0,8835.0]], ....]

当前使用的代码耗时约40-50秒，希望将时间控制在15-20秒以内：

[v.values.tolist() for k, v in df.groupby("ids")]

方法1：numpy向量化操作+分组索引 #

先将DataFrame按ids排序，确保同组数据连续，再利用numpy的数组切片替代逐组处理，大幅降低开销：

import numpy as np

# 按ids排序，保证同组数据连续
df_sorted = df.sort_values('ids')
# 转换为numpy数组，提升操作效率
arr = df_sorted.values

# 获取每个分组的起始索引
_, group_indices = np.unique(df_sorted['ids'], return_index=True)
# 按索引切割数组，得到每个分组的子数组
grouped_arrays = np.split(arr, group_indices[1:])

# 转换为目标嵌套列表格式
out = [sub_arr.tolist() for sub_arr in grouped_arrays]

该方法通过numpy的高效底层操作，避免了Pandas分组后逐组转换的额外开销，实测能将耗时压缩到目标区间内。

方法2：Pandas agg聚合优化 #

如果不想依赖numpy，可通过agg直接聚合每行数据，减少分组后的循环开销：

# 按ids分组，先将每列转为列表，再转置为行列表
out = df.groupby('ids').agg(list).apply(lambda row: list(zip(*row.values())), axis=1).tolist()

此方法利用Pandas的内置聚合函数替代手动循环，效率比原方法提升明显。

额外优化建议 #

• 将ids字段转为整数类型：df['ids'] = df['ids'].astype(int)，减少分组时的哈希计算开销
• 若内存充足，提前将DataFrame转换为numpy数组再处理，能进一步提升速度

内容的提问来源于stack exchange，提问作者Andrea Barnabò