0

Since argmax only gives one maximum values,how can we find atleast 2 or 3 elements instead of just one.

Currently my input is in the format np.argmax(array,axis=2) which is giving only one maximum and i have to extract 2 or 3 atleast from the array which is N-dimensional

desertnaut
  • 46,107
  • 19
  • 109
  • 140

2 Answers2

1

I would try to use the function called argpartition(). To get the indices of the two largest elements, do:

import numpy as np

a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])

ind = np.argpartition(a, -2)[-2:] 

ind
Out[13]: array([5, 0], dtype=int64)

a[ind]
Out[14]: array([9, 9])
Carles Sans Fuentes
  • 2,192
  • 8
  • 22
1

Using numpy.argsort. Data from @CarlesSansFuentes.

import numpy as np

a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])

args = np.argsort(-a)[:2]

array([0, 5], dtype=int64)
jpp
  • 134,728
  • 29
  • 196
  • 240