Pandas中替代for循环实现按ID分组生成排序后的level序列

问题描述 #

我有一个包含id、sort_seq和level三列的数据集，需要按id分组，依据sort_seq升序排序后提取对应的level序列。目前用for循环结合字典和列表追加的方式实现，但耗时较长，求更高效的替代方案。

输入数据集代码 #

import pandas as pd
import numpy as np
data = {'id': [1, 1, 1, 1,2, 2, 3, 3, 3, 3, 4, 5, 5, 6],
        'sort_seq': [89, 24, 56,  8,  5, 64, 93, 88, 61, 31, 50, 75,  1, 81],
        'level':['a', 'a',  'b', 'c', 'x', 'x', 'g', 'a', 'b', 'b', 'b', 'c', 'c','b']}
df = pd.DataFrame(data)

预期输出 #

按ID分组后，每个ID对应按sort_seq升序排列的level序列（如ID1对应['c','a','a','b']）

已尝试的代码 #

collect = []
for ij in df.id.unique():
  idict = {}
  x =  df[df['id'] == ij]
  x = x.sort_values(by='sort_seq',ascending=True)
  x = x.reset_index()
  idict[ij] =  x['level'].tolist()
  collect.append(idict)
collect

优化方案：利用Pandas分组聚合+排序的向量化操作 #

直接使用Pandas的groupby结合链式操作，先全局按id和sort_seq排序，再分组提取level序列，最后转为字典，效率远高于Python循环。

代码示例 #

# 先按id分组，每组内按sort_seq升序排序后提取level转为列表
result = df.sort_values(['id', 'sort_seq'], ascending=[True, True]) \
           .groupby('id')['level'] \
           .apply(list) \
           .to_dict()
print(result)

输出结果 #

{1: ['c', 'a', 'a', 'b'],
 2: ['x', 'x'],
 3: ['b', 'b', 'a', 'g'],
 4: ['b'],
 5: ['c', 'c'],
 6: ['b']}

方案优势 #

• 效率提升明显：Pandas的排序和分组操作都是底层优化的向量化实现，避免了Python循环的逐元素处理开销，数据量越大，性能差距越显著。
• 代码简洁易读：链式操作直接表达业务逻辑，无需额外的循环和字典构建步骤。
• 结果格式更实用：直接返回以id为键、level序列为值的字典，无需再处理原代码中列表嵌套字典的结构。

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