如何用Numpy实现类似Python列表排序的N维数组排序？

我们需要对3D坐标数组进行字典序排序（先按第一个元素排序，值相同则比较第二个，以此类推），要求结果既与Python原生sorted()函数对列表的排序结果完全一致，也与meshgrid + column_stack生成的原始有序坐标顺序匹配，同时要借助Numpy的排序索引，实现对其他同第一维度数组的同步排序。

问题分析 #

原代码中使用test.argsort(axis=0)的方式错误，因为该操作是对每一列单独排序，而非按行的字典序整体排序，导致结果和Python原生排序不符。

解决方案 #

使用Numpy的np.lexsort()函数可以直接实现字典序排序，该函数接收一组排序键，排序优先级从最后一个传入的键开始。对于3D坐标（x,y,z），我们需要先按x、再按y、最后按z排序，因此传入的键顺序应为(z列, y列, x列)。

修正后的测试代码 #

import numpy as np

num = 10
x = np.linspace(1, 3, num)
y = np.linspace(4, 6, num)
z = np.linspace(7, 9, num)
xg, yg, zg = np.meshgrid(x, y, z, indexing='ij', sparse=False)
test = np.column_stack((xg.ravel(), yg.ravel(), zg.ravel()))
print('Target ordering', test)

np.random.shuffle(test)
print('\nShuffled', test)

# Python原生排序结果（作为基准）
python_sorted = sorted(test.tolist())
python_sorted = np.array(python_sorted)
print('\nPython list sorted', python_sorted)

# Numpy字典序排序，获取排序索引
sort_idx = np.lexsort((test[:, 2], test[:, 1], test[:, 0]))
numpy_sorted = test[sort_idx]
print('\nNumpy lex sorted', numpy_sorted)

# 验证结果一致
print("\n排序结果是否一致：", np.array_equal(numpy_sorted, python_sorted))

同步排序其他数组 #

如果需要用该索引对其他同第一维度的数组（比如和test行数相同的特征数组）进行同步排序，直接使用索引即可：

# 示例：同步排序另一个5维特征数组
data = np.random.rand(test.shape[0], 5)
data_sorted = data[sort_idx]

内容的提问来源于stack exchange，提问作者Simon T.