Python numpy保留一个已排序的2D数组的索引列表

Question

我有一个2D numpy数组,我想创建一个新的1D数组,如果它们按升序排序,它是第一个数组中的数字索引.对于以下数组:

A = [[1,0,2],
     [0,3,0]]

我想这就像:

B = [[1,1],[0,2],[0,0],[0,1],[1,0],[1,2]]

知道如何使用预定义函数在python中完成它吗？

谢谢

Answer 1

您可以使用argsort对flattened数组的索引进行排序,然后unravel_index将平面索引转换回坐标:

>>> i = (-a).argsort(axis=None, kind='mergesort')
>>> j = np.unravel_index(i, a.shape) 
>>> np.vstack(j).T
array([[1, 1],
       [0, 2],
       [0, 0],
       [0, 1],
       [1, 0],
       [1, 2]])

-a并且kind='mergesort'是为了以降序排列稳定的数组(以匹配您正在寻找的输出).

如果您不关心稳定排序,请将第一行替换为:

>>> i = a.argsort(axis=None)[::-1]