numpy数组列表中N个最高元素的索引[重复]

Question

我需要一个函数，它从 numpy 数组的矩阵/列表中返回 N 个最高元素的索引（行和列）。

假设我有以下矩阵：

a = [[4,2,3], 
     [5,0,3]]

我想获取 N 个最高元素的索引（行、列）列表。例如，如果我想要 4 个最高元素的索引，函数应该返回

[(1,0), (0,0), (0,2), (1,2)]

我已经尝试如下实现，但它返回一个（值，行号）列表，这是我不需要的

for i, sub_list in enumerate(a):
    max_list.append((max(sub_list), i))

Answer 1

我会变平，argsort 和 unravel_index

---

f = a.ravel().argsort()
[*zip(*np.unravel_index(f[-n:], a.shape))]

[(0, 2), (1, 2), (0, 0), (1, 0)]

---

正如 yatu 所指出的，如果您有一个更大的数组，而n 则比较小，您可以将argsort 替换为np.argpartition(a.ravel(), -n)。