Encuentra el índice de los k valores más pequeños de una matriz numpy

Para encontrar el índice del valor más pequeño, puedo usar argmin :

 import numpy as np A = np.array([1, 7, 9, 2, 0.1, 17, 17, 1.5]) print A.argmin() # 4 because A[4] = 0.1 

Pero, ¿cómo puedo encontrar los índices de los valores k-más pequeños ?

Estoy buscando algo como:

 print A.argmin(numberofvalues=3) # [4, 0, 7] because A[4] <= A[0] <= A[7] <= all other A[i] 

Nota: en mi caso de uso, A tiene entre ~ 10 000 y 100 000 valores, y solo me interesan los índices de k = 10 valores más pequeños. k nunca sera> 10.

Utilice np.argpartition . No ordena toda la matriz. Solo garantiza que el elemento kth esté en posición ordenada y que todos los elementos más pequeños se muevan antes de él. Así, los primeros k elementos serán los k elementos más pequeños.

 import numpy as np A = np.array([1, 7, 9, 2, 0.1, 17, 17, 1.5]) k = 3 idx = np.argpartition(A, k) print(idx) # [4 0 7 3 1 2 6 5] 

Esto devuelve los valores k-más pequeños. Tenga en cuenta que estos pueden no estar ordenados.

 print(A[idx[:k]]) # [ 0.1 1. 1.5] 

Para obtener los valores k-más grandes usa

 idx = np.argpartition(A, -k) # [4 0 7 3 1 2 6 5] A[idx[-k:]] # [ 9. 17. 17.] 

ADVERTENCIA: No (re) use idx = np.argpartition(A, k); A[idx[-k:]] idx = np.argpartition(A, k); A[idx[-k:]] para obtener el k-más grande. Eso no siempre funcionará. Por ejemplo, estos NO son los 3 valores más grandes en x :

 x = np.array([100, 90, 80, 70, 60, 50, 40, 30, 20, 10, 0]) idx = np.argpartition(x, 3) x[idx[-3:]] array([ 70, 80, 100]) 

Aquí hay una comparación con np.argsort , que también funciona, pero simplemente ordena la matriz completa para obtener el resultado.

 In [2]: x = np.random.randn(100000) In [3]: %timeit idx0 = np.argsort(x)[:100] 100 loops, best of 3: 8.26 ms per loop In [4]: %timeit idx1 = np.argpartition(x, 100)[:100] 1000 loops, best of 3: 721 µs per loop In [5]: np.alltrue(np.sort(np.argsort(x)[:100]) == np.sort(np.argpartition(x, 100)[:100])) Out[5]: True 

Puedes usar numpy.argsort con rebanar

 >>> import numpy as np >>> A = np.array([1, 7, 9, 2, 0.1, 17, 17, 1.5]) >>> np.argsort(A)[:3] array([4, 0, 7], dtype=int32) 

Para matrices n-dimensional , esta función funciona bien. Los indecies se devuelven en forma callable. Si desea que se devuelva una lista de los índices, debe transponer la matriz antes de hacer una lista.

Para recuperar el k más grande, simplemente pase en -k .

 def get_indices_of_k_smallest(arr, k): idx = np.argpartition(arr.ravel(), k) return tuple(np.array(np.unravel_index(idx, arr.shape))[:, range(min(k, 0), max(k, 0))]) # if you want it in a list of indices . . . # return np.array(np.unravel_index(idx, arr.shape))[:, range(k)].transpose().tolist() 

Ejemplo:

 r = np.random.RandomState(1234) arr = r.randint(1, 1000, 2 * 4 * 6).reshape(2, 4, 6) indices = get_indices_of_k_smallest(arr, 4) indices # (array([1, 0, 0, 1], dtype=int64), # array([3, 2, 0, 1], dtype=int64), # array([3, 0, 3, 3], dtype=int64)) arr[indices] # array([ 4, 31, 54, 77]) %%timeit get_indices_of_k_smallest(arr, 4) # 17.1 µs ± 651 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) 

numpy.partition(your_array, k) es una alternativa. No es necesario rebanar ya que proporciona los valores ordenados hasta el elemento kth .