如何优化Pandas DataFrame按ID分组的迭代代码？

问题背景 #

你需要将DataFrame按ID分组，把每组的start、end、label字段整合成子列表，最终输出指定格式的列表。原始数据和期望结果如下：

原始数据 #

import pandas as pd

data = {
  "ID": ["1", "1", "2"],
  "start": [5, 6, 30],
  "end": [10,20,50],
  "label": ["age", "gender", "history"]
}

df = pd.DataFrame(data)

期望输出 #

[{'ID': 1, 'Labels': [[5,10,'age'], [6,20,"gender"]]}, {'ID': 2, 'Labels': [[30,50,'history']]} ]

你的原始代码采用双重循环，不仅时间复杂度高（O(n²)），还存在逻辑错误，导致结果不符合预期且运行缓慢。

原始代码的问题 #

1. 双重循环导致性能爆炸：遍历每个索引时又重新遍历所有行，数据量稍大就会耗时剧增。
2. 逻辑错误：每次外层循环都会重复添加同分组的内容，且在ID不匹配时错误追加字典，最终结果会有大量冗余数据。

优化方案 #

方案一：用pandas内置groupby实现（最优性能） #

利用pandas的groupby结合聚合操作，内部是向量化实现，性能远高于手动循环：

# 先把ID转为整数（匹配期望结果的格式）
df['ID'] = df['ID'].astype(int)

# 按ID分组，将每组的start/end/label转为子列表，再转为目标格式
result = (
    df.groupby('ID')
      .apply(lambda group: group[['start', 'end', 'label']].values.tolist())
      .reset_index(name='Labels')
      .to_dict('records')
)

print(result)

方案二：单次遍历手动分组（直观易懂） #

如果不想依赖groupby，也可以用字典暂存分组数据，只遍历DataFrame一次：

group_cache = {}
# 遍历每一行
for _, row in df.iterrows():
    id_num = int(row['ID'])
    # 组装当前行的标签列表
    label_item = [row['start'], row['end'], row['label']]
    # 缓存中没有当前ID则初始化列表，否则追加
    if id_num not in group_cache:
        group_cache[id_num] = []
    group_cache[id_num].append(label_item)

# 转换为期望的列表格式
result = [{'ID': id_val, 'Labels': labels} for id_val, labels in group_cache.items()]

print(result)

效果说明 #

两种方案的时间复杂度都是O(n)，相比原代码的O(n²)，数据量越大性能提升越显著。其中方案一利用pandas的优化实现，性能最优；方案二更适合理解分组逻辑。

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