I need a function that returns the index (row and column) of the N highest element from a matrix/list of numpy array.
Suppose that I have the following matrix:
a = [[4,2,3],
[5,0,3]]
I would like to get a list of the indexes (row,column) of the N highest elements. For example, if I want the index of the 4 highest element, the function should return
[(1,0), (0,0), (0,2), (1,2)]
I have tried implementing this as following, but it returns a list of (value, row number), which is not want I need
for i, sub_list in enumerate(a):
max_list.append((max(sub_list), i))
I would flatten, argsort
, and unravel_index
f = a.ravel().argsort()
[*zip(*np.unravel_index(f[-n:], a.shape))]
[(0, 2), (1, 2), (0, 0), (1, 0)]
As yatu points out, if you have a larger array, and a small-ish n
, you can replace argsort
with np.argpartition(a.ravel(), -n)
.
Collected from the Internet
Please contact [email protected] to delete if infringement.
Comments