按多个轴排序二维numpy数组

Question

我有一个形状为(N,2)的二维numpy数组，这个数组里存放着N个点的坐标（x和y）。比如说：

array([[3, 2],
       [6, 2],
       [3, 6],
       [3, 4],
       [5, 3]])

我想把这些点按照x坐标的大小进行排序，如果x坐标相同的话，再按照y坐标来排序。所以上面的数组应该变成这样：

array([[3, 2],
       [3, 4],
       [3, 6],
       [5, 3],
       [6, 2]])

如果这是一个普通的Python列表，我可以简单地定义一个比较函数来实现我的需求，但据我所知，numpy的排序函数不支持用户自定义的比较函数。有没有什么好的办法呢？

---

补充：谢谢大家的建议！我做了一个快速测试，生成了1000000个随机整数点，并对我能运行的排序方法进行了基准测试（抱歉，现在不能升级numpy）。

Mine:   4.078 secs 
mtrw:   7.046 secs
unutbu: 0.453 secs

Answer 1

这个 numpy_indexed 包（声明：我是它的作者）可以用来高效地处理多维数组的问题，完全使用向量化的方式来解决。

import numpy_indexed as npi
npi.sort(a)  # by default along axis=0, but configurable

Answer 2

标题提到的是“对二维数组进行排序”。虽然提问者使用的是一个形状为(N,2)的数组，但其实可以把unutbu的解决方案推广到任何形状为(N,M)的数组，因为这可能是大家真正想要的。

我们可以先转置这个数组，然后用负的步长来使用切片语法，这样就可以把所有的列以相反的顺序传递给lexsort进行排序：

>>> import numpy as np
>>> a = np.random.randint(1, 6, (10, 3))
>>> a
array([[4, 2, 3],
       [4, 2, 5],
       [3, 5, 5],
       [1, 5, 5],
       [3, 2, 1],
       [5, 2, 2],
       [3, 2, 3],
       [4, 3, 4],
       [3, 4, 1],
       [5, 3, 4]])

>>> a[np.lexsort(np.transpose(a)[::-1])]
array([[1, 5, 5],
       [3, 2, 1],
       [3, 2, 3],
       [3, 4, 1],
       [3, 5, 5],
       [4, 2, 3],
       [4, 2, 5],
       [4, 3, 4],
       [5, 2, 2],
       [5, 3, 4]])

Answer 3

使用 lexsort：

import numpy as np    
a = np.array([(3, 2), (6, 2), (3, 6), (3, 4), (5, 3)])

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

a[ind]
# array([[3, 2],
#       [3, 4],
#       [3, 6],
#       [5, 3],
#       [6, 2]])

---

a.ravel() 这个方法会返回一个视图，如果 a 是 C_CONTIGUOUS 的话。简单来说，如果这个条件成立，@ars的方法，稍微改动一下，用 ravel 替代 flatten，就能很好地在原地对 a 进行排序：

a = np.array([(3, 2), (6, 2), (3, 6), (3, 4), (5, 3)])
dt = [('col1', a.dtype),('col2', a.dtype)]
assert a.flags['C_CONTIGUOUS']
b = a.ravel().view(dt)
b.sort(order=['col1','col2'])

因为 b 是 a 的一个视图，所以排序 b 的时候，a 也会被排序：

print(a)
# [[3 2]
#  [3 4]
#  [3 6]
#  [5 3]
#  [6 2]]