CODEWITHSUNDEEP
pythonarraysnumpy
rachana_sharma003
rachana_sharma003

Reputation: 127

How to get N maximum values in a multi dimensional numpy array along a given axis(say 2)?

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

Upvotes: 0

Views: 611

Answers (2)

jpp
jpp

Reputation: 164823

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)

Upvotes: 1

Carles
Carles

Reputation: 2829

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

Upvotes: 1

Related Questions