0

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))
user3483203
  • 45,503
  • 8
  • 43
  • 75
Message Passing
  • 107
  • 2
  • 12
  • I think you have an error in your output. `(1, 2)` should be included instead of `(0, 1)`, since 3 > 2 – user3483203 Feb 14 '20 at 14:29
  • Yes you are right. I have modified the question, thanks! – Message Passing Feb 14 '20 at 14:36
  • 1
    Does this answer your question? [How do I get indices of N maximum values in a NumPy array?](https://stackoverflow.com/questions/6910641/how-do-i-get-indices-of-n-maximum-values-in-a-numpy-array) There are solutions for 2D arrays as well. – Georgy Feb 14 '20 at 15:02

1 Answers1

4

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).

user3483203
  • 45,503
  • 8
  • 43
  • 75