按多个轴对2D 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的sort函数不接受用户自定义的比较器。有什么想法吗？

---

编辑：感谢提供的意见！我设置了一个快速测试案例，包含1000000个随机整数点，并对我能运行的进行了基准测试（抱歉，目前无法升级numpy）。

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

- perimosocordiae

Answer 1

使用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 是 C_CONTIGUOUS 的，a.ravel() 返回一个视图。如果这是真的，那么稍作修改 @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]]

- unutbu

---

好的，我会尽力以简洁明了的方式翻译以下内容，同时保持原意不变：是的，我经常很难理解文档。例子往往更能让人理解。但问题在于，当我跟着例子操作后，重新阅读文档时发现文档本来就非常清晰易懂... :-) - unutbu

@Noah：是的，这将创建一个新数组。 - unutbu

@Noah：我已经修改了我的答案，展示了如何在numpy数组上按多个索引进行原地排序。 - unutbu

请注意，lexsort 使用序列中的最后一个条目作为主键，倒数第二个作为次要键等。这让我措手不及。这个答案做得很对，但很容易被忽视。 - LucasB

对于任意的二维数组：np.lexsort(a.T[::-1]) - scleronomic

Answer 2

标题说“对2D数组进行排序”。虽然提问者使用了一个形状为(N,2)的数组，但可以将unutbu的解决方案推广到任何(N,M)数组上，因为这可能是人们实际上正在寻找的。

可以使用transpose函数对数组进行转置，并使用负step的切片符号将所有列以相反的顺序传递给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]])

- arekolek

---

如果我只想按列1排序，然后按列3排序。我该如何更改您的代码？ - BlueJapan

Answer 3

numpy_indexed包（免责声明：我是它的作者）可以以高效的完全向量化方式解决这些nd数组上的处理问题：

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

- Eelco Hoogendoorn

Answer 4

我曾经也遇到过同样的问题，但是现在已经得到了帮助并解决了这个问题。如果你的数组有列名（结构化数组），那么它可以顺利地工作。我认为这是一种非常简单的方法，使用与Excel相同的逻辑进行排序：

array_name[array_name[['colname1','colname2']].argsort()]

请注意双括号包含排序条件。当然，您可以使用多于2列作为排序条件。

- JaimeBee

Answer 5

你可以使用 np.complex_sort。这会将你的数据转换为浮点数，希望这不是问题：

>>> a = np.array([[3, 2], [6, 2], [3, 6], [3, 4], [5, 3]])
>>> atmp = np.sort_complex(a[:,0] + a[:,1]*1j)
>>> b = np.array([[np.real(x), np.imag(x)] for x in atmp])
>>> b
array([[ 3.,  2.],
       [ 3.,  4.],
       [ 3.,  6.],
       [ 5.,  3.],
       [ 6.,  2.]])

- mtrw

---

我认为你赢得了聪明才智奖；我从未想过将y坐标设为虚数！ - perimosocordiae

但是速度太慢了！抱歉，我在发布时没有真正考虑性能。 - mtrw

Answer 6

编辑：已删除错误答案。

以下是一种使用中间结构化数组的方法：

from numpy import array

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

b = a.flatten()
b.dtype = [('x', '<i4'), ('y', '<i4')]
b.sort()
b.dtype = '<i4'
b.shape = a.shape

print b

这将会输出所需的结果：

[[3 2]
 [3 4]
 [3 6]
 [5 3]
 [6 2]]

不确定这是否是最好的方法。

- ars

---

这并不完全可行，因为它会失去我的点之间x和y的关联。 - perimosocordiae

嗯，当我运行它时，在b.shape = a.shape这一行上出现了一个错误：“ValueError：新数组的总大小必须保持不变”。我正在运行Python 2.6.2，带有numpy 1.2.1。 - perimosocordiae

我正在运行Python 2.5.4和numpy 1.3.0。尝试升级numpy的版本。 - ars

Answer 7

我找到了一种方法来做到这一点：

from numpy import array
a = array([(3,2),(6,2),(3,6),(3,4),(5,3)])
array(sorted(sorted(a,key=lambda e:e[1]),key=lambda e:e[0]))

需要两次排序（而且使用普通的 Python sorted 函数而不是更快的 numpy sort）确实很糟糕，但是它可以很好地适合一行。

- perimosocordiae