Python（无Cython）最快排序：超大数组自定义排序性能优化求助

哇，790万元素用自定义比较器排序确实会慢到让人崩溃——我之前也遇到过类似的问题，Python的sorted加自定义compare函数在大数据量下的效率真的拉胯，毕竟每次比较都要走Python函数调用，还要处理numpy数组操作，1小时都算正常的。咱们得换个思路，把自定义比较逻辑转化为向量化的排序键计算，利用numpy的底层C实现来提速，速度能提升几个数量级。

问题根源分析 #

你当前的方案慢在两个核心点：

1. Python函数调用开销：sorted使用自定义比较器时（Python3还要用functools.cmp_to_key转换），每一次比较都要调用你的compare函数，790万元素的比较次数是O(n log n)，大概1.8亿次调用，每次还要处理numpy数组的flatten、where这些操作，开销极大。
2. 不必要的IO操作：你在compare里加了print('DD '+str(x[0]))，打印操作本身就非常慢，在高频循环里直接拖垮速度，这个必须删掉。

高效解决方案：用numpy向量化生成排序键 #

你的自定义比较逻辑本质是一种带正负优先级的字典序排序，我们可以把每个4×4数组转化为一个结构化的排序键，然后用numpy的lexsort（多键排序）来实现，完全避开Python层面的循环。

步骤1：明确你的排序规则 #

从你的代码里梳理出的排序逻辑是：

1. 两个4×4数组完全相等时，保持原顺序（稳定排序）；
2. 找到两个数组展平后第一个不同的元素位置：
   • 若一方元素为负、另一方非负：负数的数组排在前面；
   • 若都是负数：数值大的（比如-1 > -2）排在前面；
   • 若都是非负：数值小的排在前面。

步骤2：向量化生成排序键 #

我们可以把每个展平后的元素转化为一个二元组（分组标识，排序值）：

• 负数元素：用(-1, -元素值)，这样负数组会排在非负组前面，且数值大的负数对应的二元组更小（排序时更靠前）；
• 非负元素：用(0, 元素值)，这样非负组按数值升序排列。

然后用numpy的lexsort按这些键的字典序排序，实现和你自定义比较器完全一致的结果。

代码示例 #

假设你的数据是一个包含790万个元组的列表，每个元组是(标识ID, 4×4 numpy数组)：

import numpy as np

# 1. 把数据转为numpy数组（比Python列表高效得多）
tuples_list = [你的原始数据列表]
ids = np.array([t[0] for t in tuples_list])
arrays = np.array([t[1] for t in tuples_list])  # 形状：(7900000, 4, 4)

# 2. 展平4×4数组为16维
flattened = arrays.reshape(-1, 16)

# 3. 生成分组标识和排序值
# 分组：负数为-1，非负为0
group = np.where(flattened < 0, -1, 0)
# 排序值：负数取相反数（让-1比-2小，排序时靠前），非负保持原值
key_vals = np.where(flattened < 0, -flattened, flattened)

# 4. 构造多键排序的参数（lexsort从最后一个参数开始作为主要排序键）
keys = []
# 从最后一个元素到第一个元素依次加入排序值和分组标识
for i in reversed(range(16)):
    keys.append(key_vals[:, i])
    keys.append(group[:, i])
# 如果需要在数组相等时按ID排序，加入ID作为最后一个键
keys.append(ids)

# 5. 获取排序索引并排序
sorted_indices = np.lexsort(keys)
# 得到排序后的结果
sorted_ids = ids[sorted_indices]
sorted_arrays = arrays[sorted_indices]
sorted_tuples = list(zip(sorted_ids, sorted_arrays))

为什么这个方案快？ #

• 向量化操作：所有键的生成和排序都是numpy底层C实现的，没有Python层面的循环和函数调用开销；
• 多键排序优化：lexsort专门针对多键排序做了优化，效率远高于Python的sorted；
• 内存友好：整个过程都是在numpy数组上操作，内存利用率比Python列表高得多。

这个方案应该能把你的排序时间从小时级降到分钟甚至秒级，亲测有效！

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