Python numpy 保留已排序二维数组的索引列表

Posted 2023-02-23

【中文标题】Python numpy 保留已排序二维数组的索引列表

【英文标题】：Python numpy keep a list of indices of a sorted 2D array

【发布时间】：2015-06-26 09:32:54

【问题描述】：

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

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

我希望是这样的：

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

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

谢谢

【参考方案1】：

您可以使用argsort 对扁平数组的索引进行排序，然后使用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]

【讨论】：

关于使用np.unravel_index 的提示非常棒！我很高兴学习这种方法。

“稳定排序”是什么意思？这是否会对相同的值产生一致的排序？

@stvn66 我相信这确实会为相同的值产生一致的排序，这等于稳定的排序