Trouver l'indice des k plus petites valeurs d'un tableau de numpy
pour trouver l'indice de la plus petite valeur, je peux utiliser 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
Mais comment puis-je trouver les indices des k-plus petites valeurs ?
je cherche quelque chose comme:
print A.argmin(numberofvalues=3)
# [4, 0, 7] because A[4] <= A[0] <= A[7] <= all other A[i]
Note: dans mon cas D'utilisation A A entre ~ 10 000 et 100 000 valeurs, et je ne m'intéresse que pour les indices des plus petites valeurs de k=10. K ne sera jamais > 10.
4 réponses
Utilisation np.argpartition . Il ne trie pas le tableau entier. Il garantit seulement que l'élément kth est en position triée et que tous les éléments plus petits seront déplacés avant lui. Ainsi, les premiers éléments k seront les éléments k-les plus petits.
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]
renvoie les valeurs k-les plus faibles. Notez qu'ils peuvent ne pas être dans l'ordre trié.
print(A[idx[:k]])
# [ 0.1 1. 1.5]
pour obtenir le k plus grandes valeurs d'usage
idx = np.argpartition(A, -k)
# [4 0 7 3 1 2 6 5]
A[idx[-3:]]
# [ 9. 17. 17.]
Voici une comparaison avec np.argsort , qui fonctionne aussi mais trie juste l'ensemble du tableau pour obtenir le résultat.
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
vous pouvez utiliser numpy.argsort avec tranchage
>>> 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)
numpy.partition(your_array, k) est une alternative. Aucun tranchage nécessaire car il donne les valeurs triées jusqu'à l'élément kth .
pour n-dimensional arrays , cette fonction fonctionne bien. Les indecies sont retournées sous une forme sujette à appel. Si vous voulez qu'une liste des indices soit retournée, vous devez transposer le tableau avant de faire une liste.
pour récupérer le k le plus grand, il suffit de passer dans -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()
exemple:
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)