python - 返回对应于排序的 ndarray 的所有元素的索引的元组列表的函数？

Question

我知道 numpy.argsort()，但它的作用是返回数组中将沿特定轴排序的元素的索引。

我需要的是对 N 维数组中的所有值进行排序，并得到一个元组的线性列表作为结果。

像这样:

>>> import numpy
>>> A = numpy.array([[7, 8], [9, 5]])
>>> numpy.magic(A)
[(1, 0), (0, 1), (0, 0), (1, 1)]

附言对于这个数组，我什至不明白 argsort 的输出试图告诉我什么。

Answer 1

np.argsort(A) 分别对 A 的每一行进行排序。例如，

In [21]: np.argsort([[6,5,4],[3,2,1]])
Out[21]: 
array([[2, 1, 0],
       [2, 1, 0]])

相反，您想将数组展平为一维值数组，然后对其进行 argsort。这可以通过将 axis 参数设置为 None 来完成(感谢@Akavall 指出这一点):

In [23]: np.argsort(A, axis=None)
Out[23]: array([3, 0, 1, 2])

然后使用np.unravel_index恢复 A 中的关联索引。

---

In [14]: import numpy as np

In [15]: A = np.array([[7, 8], [9, 5]])   

In [4]: np.column_stack(np.unravel_index(np.argsort(A, axis=None)[::-1], A.shape))
Out[4]: 
array([[1, 0],
       [0, 1],
       [0, 0],
       [1, 1]])

请注意，对于 NumPy 版本 1.5.1 或更早版本，np.unravel_index 如果为其第一个参数传递类似数组的对象，则会引发 ValueError。在这种情况下，您可以使用列表理解:

In [17]: [np.unravel_index(p, A.shape) for p in np.argsort(A, axis=None)[::-1]]
Out[17]: [(1, 0), (0, 1), (0, 0), (1, 1)]