如何提升大数据集下Python分组代码的运行效率？

慢因确认 #

你的代码运行缓慢完全是写法低效导致的，不是单纯数据量大的问题，核心有两个致命缺陷：

• 双重冗余遍历：外层遍历288个唯一order，内层每次都要完整遍历159万条全量数据，总遍历次数超过4.6亿次，时间开销直接拉满。
• 列表线性查询：每次用picker元素 in hold判断是否重复，列表的in操作是线性扫描，每次查询都是O(n)复杂度，进一步拖慢速度。

优化思路 #

核心是一次遍历完成分组+去重，同时提前记录picker首次出现的顺序，最后按要求输出格式整理：

1. 先遍历一次picker列表，记录每个picker首次出现的索引，用来保证后续子列表内的picker顺序符合要求。
2. 用字典做分组容器，key为order值，value用集合存储对应的picker（集合的in判断是O(1)，自动去重），一次遍历就能完成所有分组。
3. 最后将字典的key按从小到大排序，每个key对应的picker按首次出现索引排序，生成目标格式的pairs列表。

优化后代码 #

order = [1, 2, 3, 4, 1, 5, 3, 6, 7, 1, 8, 9, 4, 4, 2, 8, 4, 4, 2]
picker = ['a', 'b', 'c', 'd', 'a', 'e', 'c', 'f', 'g', 'a', 'h', 'i', 'j', 'k', 'b', 'h', 'j', 'j', 'k']

# 记录每个picker首次出现的索引
first_occur = {}
for idx, p in enumerate(picker):
    if p not in first_occur:
        first_occur[p] = idx

# 一次遍历完成order与picker的分组去重
order_groups = {}
for o, p in zip(order, picker):
    if o not in order_groups:
        order_groups[o] = set()
    order_groups[o].add(p)

# 生成符合要求的pairs列表
pairs = []
# 按order值从小到大排序遍历
for o in sorted(order_groups.keys()):
    # 将当前order对应的picker按首次出现索引排序
    sorted_pickers = sorted(order_groups[o], key=lambda x: first_occur[x])
    # 组合成目标子列表格式
    pairs.append([o] + sorted_pickers)

print(pairs)

优化效果 #

• 时间复杂度从原代码的O(M*N)（M为唯一order数，N为总数据量）降至O(N + K*L log L)（N为总数据量，K为唯一order数，L为单order对应的picker数），针对159万条数据，运行速度会提升数百倍甚至上千倍。
• 输出格式完全符合要求：pairs子列表按order值从小到大排列，子列表内的picker元素按其首次出现的顺序排列。

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