¿Cómo obtengo índices de N valores máximos en una matriz NumPy?
NumPy propone una forma de obtener el índice del valor máximo de una matriz mediante np.argmax.
Me gustaría algo similar, pero devolviendo los índices de los Nvalores máximos.
Por ejemplo, si tengo una matriz [1, 3, 2, 4, 5], nargmax(array, n=3)devolvería los índices [4, 3, 1]que corresponden a los elementos [5, 4, 3].
Las versiones más nuevas de NumPy (1.8 y superiores) tienen una función llamada argpartitionpara esto. Para obtener los índices de los cuatro elementos más grandes, haga
>>> a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> a
array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> ind = np.argpartition(a, -4)[-4:]
>>> ind
array([1, 5, 8, 0])
>>> top4 = a[ind]
>>> top4
array([4, 9, 6, 9])
A diferencia de argsort, esta función se ejecuta en tiempo lineal en el peor de los casos, pero los índices devueltos no están ordenados, como se puede ver en el resultado de la evaluación a[ind]. Si también lo necesitas, ordénalos después:
>>> ind[np.argsort(a[ind])]
array([1, 8, 5, 0])
Para ordenar los k elementos superiores de esta manera se necesita O( n + k log k ) tiempo.
Lo más simple que se me ha ocurrido es:
>>> import numpy as np
>>> arr = np.array([1, 3, 2, 4, 5])
>>> arr.argsort()[-3:][::-1]
array([4, 3, 1])
Esto implica una clasificación completa de la matriz. Me pregunto si numpyproporciona una forma integrada de realizar una clasificación parcial; hasta ahora no he podido encontrar ninguno.
Si esta solución resulta demasiado lenta (especialmente para archivos pequeños n), puede que valga la pena intentar codificar algo en Cython .
Más simple aún:
idx = (-arr).argsort()[:n]
donde n es el número de valores máximos.
