NumPy lexsort二级排序结果不符预期的原因及解决方法

核心问题：np.lexsort()的参数顺序搞反了！ #

numpy的np.lexsort()函数有个很容易踩坑的特性：它的参数是从「最次要排序键」到「最主要排序键」排列的——简单说，最后传入的数组才是第一优先级的排序关键字（主关键字），前面的参数都是用来处理主关键字值相同的情况的次要键。

结合你的例子分析 #

你执行的代码是：

i = np.lexsort((a[:,0], a[:,1]))

这相当于告诉numpy：

1. 优先按**第二列（a[:,1]）**升序排序（主关键字）
2. 只有当第二列的值完全相同时，才用第一列（a[:,0]）来排序

这就是为什么你的结果里，第二列是200→300→400→1000的升序，而第一列的两个1被拆开了——因为其中一个的第二列值是1000，比300、400都大，所以排在了最后。

实现预期结果的方法 #

如果你想要先按第一列升序，第一列相同的情况下再按第二列升序，只需要把参数顺序反过来，把主关键字（第一列）放在最后：

import numpy as np

# 你的原始数组
a = np.array([ [1, 1000, 58], [1, 200, 69], [3, 300, 34], [4, 400, 82]], dtype = int)

# 次要键（第二列）在前，主键（第一列）在后
i = np.lexsort((a[:,1], a[:,0]))
# 按排序索引提取结果
expected_result = a[i]

print(expected_result)

运行后会输出你预期的结果：

array([[   1,  200,   69],
       [   1, 1000,   58],
       [   3,  300,   34],
       [   4,  400,   82]])

另一种可选方法：结构化数组排序 #

如果你觉得lexsort()的参数顺序容易搞混，也可以把数组转换成结构化数组，用np.sort()的order参数明确指定排序顺序，可读性更强：

# 将数组转换为结构化数组，指定列名
a_struct = a.view('i,i,i').reshape(-1)
a_struct.dtype.names = ('col0', 'col1', 'col2')
# 按col0（第一列）、col1（第二列）排序
sorted_struct = np.sort(a_struct, order=['col0', 'col1'])
# 转换回普通数组格式
expected_result = sorted_struct.view(int).reshape(-1, 3)

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