如何找到 n 维 numpy 数组的第 i 个最大元素的索引？

Question

我知道如何找到 n 维数组的最大元素的索引。

我们举个例子：

a=np.asarray([[1,7,-4],[9,-11,-17]])

然后（source）：

from numpy import unravel_index
unravel_index(a.argmax(), a.shape)

返回：

(1, 0)

确实a[1,0] 是9，这是数组a 中的最高元素，所以我们很好。

---

我还能够弄清楚如何找到一维 numpy 数组 (using) 的第 i 个最大元素的索引：

a = np.array([1, 3, 2, 4, 5])

i=3 # we want the third largest element, for example
a.argsort()[-i]

这会返回1，这很好，因为a[1]=3 确实是a 的第三大元素。

---

我想将这两者结合起来。所以如果我有

a=np.asarray([[1,7,-4],[9,-11,-17]])

我想要一个输出告诉我数组a 的第i 个最大元素的索引，例如如果i=3，输出应该是[0,0]，因为a[0,0]=1 是第i 个（第三个) a 的最大元素。

我该怎么做？

Answer 1

那么要得到某个最大的索引或者哪个，你可以使用where：

通过webDev添加到上述答案：

import numpy as np
i=2

a=np.asarray([[1,7,-4],[9,-11,-17]])

flat=a.flatten()
flat.sort()
tryvalue= flat[-i]

i, j = np.where(a == tryvalue)
print(i,j)

这会给你：

[0] [1]

我的意思是您可以自行更改您希望这些索引的样子（元组或其他）。

Answer 2

这是一种简单的方法。

import numpy as np
i=3
a=np.asarray([[1,7,-4],[9,-11,-17]])
flat=a.flatten()
flat.sort()
print(flat)
print(flat[-i])
i, j = np.where(a == flat[-i])
print(i,j)

您可以将其展平然后排序。它将根据您的第 i 个最大输出，即i=3，为您提供所需的输出。如果你输入 i=5，flat[-i] 会给你-11。

Answer 3

您还可以在展平数组上使用 heapq.nlargest 并获取那些最大的 i 元素中的最小值。在大多数情况下，这应该比排序更有效：

import numpy as np
import heapq

a = np.asarray([[1, 7, -4], [9, -11, -17]])
i = 2

ith_largest = min(heapq.nlargest(i, a.flatten()))
x, y = np.where(a == ith_largest)
print(x, y)  # [0] [1]

Answer 4

Amit Amola 的answer 非常好。如果有人需要另一个，我找到了这个解决方案：

a=np.asarray([[1,7,-4],[9,-11,-17]])
flat=a.flatten()
flat.sort()
i=5

for k, p in enumerate(a):
    for j, q in enumerate(p):
        if q == flat[-i]:
            indices=[k,j]

print(indices)

给[1, 1]，很好。