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

Question

I know how to find the indices of the maximum element of an n-dimensional array.

Let's have for example:

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

Then ( source ):

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

returning:

(1, 0)

and indeed a[1,0] is 9 , which is the highest element in the array a , so we are good.

---

I am also able to figure out how to find the indices of the ith largest element of a one-dimensional numpy array ( using ):

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

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

This returns 1 , which is good, since a[1]=3 which is indeed the third largest element of a .

---

I would like to combine these two. So if I have

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

I would like to get an output telling me the indices of the ith largest element of the array a , for example if i=3 , the output should be [0,0] , since a[0,0]=1 is the ith (third) largest element of a .

How can I do this?

Answer 1

Well to get the index of some largest or whichever, you can use where:

Adding to above answer by 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)

This will give you:

[0] [1]

I mean you can make changes on your own that how you want these indexes to be like(tuple or whatever).

Answer 2

This is a simple way to do so.

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)

You can flatten and then sort it. It will give you the output that you want based on your ith largest ie i=3 . If you enter i=5, flat[-i] will give you -11 .

Answer 3

You can also use heapq.nlargest on the flattened array and get the minimum of those largest i elements. That should be more efficient than sorting for most cases:

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 's answer is perfectly good. In case someone needs another one, I found this solution:

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)

Giving [1, 1] , which is good.