0

For each row, I would like to get the index of the largest two values.

import numpy as np

x = np.array([[7, 5, 6],
              [4, 9, 3],
              [1, 6, 7]])

Here's the result I would like to obtain:

[0, 2]
[1, 0]
[2, 1]
Nicolas Gervais
  • 33,817
  • 13
  • 115
  • 143
  • 4
    You forgot to post your attempt to solve this problem. – Scott Hunter Jun 30 '22 at 16:14
  • Are the sublists always of length 3? If so, get the index of the minimum and use that to eliminate the index that isn't wanted. – ChaddRobertson Jun 30 '22 at 16:24
  • Additionally, does the order of the resulting indices matter? Your desired output seems to prioritise the largest value of the array first. – ChaddRobertson Jun 30 '22 at 16:27
  • 2
    Does this answer your question? [Is it possible to use argsort in descending order?](https://stackoverflow.com/questions/16486252/is-it-possible-to-use-argsort-in-descending-order) – Stef Jun 30 '22 at 16:27
  • They are not of length 3, it's a minimal example. And yes, the order of the resulting indices matters. – Nicolas Gervais Jun 30 '22 at 16:44

1 Answers1

1

Use np.argsort along the second axis and take last two values reversed.

import numpy as np

x = np.array([[7, 5, 6],
              [4, 9, 3],
              [1, 6, 7]])

ind = x.argsort(axis=1)[:,-1:-3:-1]
print(ind)
[[0 2]
 [1 0]
 [2 1]]
AboAmmar
  • 5,439
  • 2
  • 13
  • 24