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

Posted 2023-02-23

【中文标题】numpy数组列表中N个最高元素的索引[重复]

【英文标题】：index of N highest elements from a list of numpy array [duplicate]

【发布时间】：2020-05-30 08:11:15

【问题描述】：

我需要一个函数，它从 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))

【问题讨论】：

我认为您的输出中有错误。应该包括(1, 2) 而不是(0, 1)，因为 3 > 2

是的，你是对的。我已经修改了问题，谢谢！

这能回答你的问题吗？ How do I get indices of N maximum values in a NumPy array? 二维数组也有解决方案。

【参考方案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)。