Numpy : trouver l'index des éléments dans la plage.

Question

J'ai un tableau numpy de nombres, par exemple,

a = np.array([1, 3, 5, 6, 9, 10, 14, 15, 56])  

Je voudrais trouver tous les indices des éléments dans une plage spécifique. Par exemple, si l'intervalle est (6, 10), la réponse devrait être (3, 4, 5). Existe-t-il une fonction intégrée permettant de faire cela ?

Answer 1

Vous pouvez utiliser np.where pour obtenir des indices et np.logical_and pour fixer deux conditions :

import numpy as np
a = np.array([1, 3, 5, 6, 9, 10, 14, 15, 56])

np.where(np.logical_and(a>=6, a<=10))
# returns (array([3, 4, 5]),)

Answer 2

Comme dans la réponse de @deinonychusaur, mais encore plus compact :

In [7]: np.where((a >= 6) & (a <=10))
Out[7]: (array([3, 4, 5]),)

Answer 3

Résumé des réponses

Pour comprendre quelle est la meilleure réponse, nous pouvons effectuer un chronométrage en utilisant les différentes solutions. Malheureusement, la question n'était pas bien posée donc il y a des réponses à des questions différentes, ici j'essaie d'indiquer la réponse à la même question. Étant donné le tableau :

a = np.array([1, 3, 5, 6, 9, 10, 14, 15, 56])

La réponse devrait être la suivante indices des éléments entre une certaine fourchette, que nous supposons inclusive, dans ce cas, 6 et 10.

answer = (3, 4, 5)

Correspondant aux valeurs 6,9,10.

Pour tester la meilleure réponse, nous pouvons utiliser ce code.

import timeit
setup = """
import numpy as np
import numexpr as ne

a = np.array([1, 3, 5, 6, 9, 10, 14, 15, 56])
# or test it with an array of the similar size
# a = np.random.rand(100)*23 # change the number to the an estimate of your array size.

# we define the left and right limit
ll = 6
rl = 10

def sorted_slice(a,l,r):
    start = np.searchsorted(a, l, 'left')
    end = np.searchsorted(a, r, 'right')
    return np.arange(start,end)
"""

functions = ['sorted_slice(a,ll,rl)', # works only for sorted values
'np.where(np.logical_and(a>=ll, a<=rl))[0]',
'np.where((a >= ll) & (a <=rl))[0]',
'np.where((a>=ll)*(a<=rl))[0]',
'np.where(np.vectorize(lambda x: ll <= x <= rl)(a))[0]',
'np.argwhere((a>=ll) & (a<=rl)).T[0]', # we traspose for getting a single row
'np.where(ne.evaluate("(ll <= a) & (a <= rl)"))[0]',]

functions2 = [
   'a[np.logical_and(a>=ll, a<=rl)]',
   'a[(a>=ll) & (a<=rl)]',
   'a[(a>=ll)*(a<=rl)]',
   'a[np.vectorize(lambda x: ll <= x <= rl)(a)]',
   'a[ne.evaluate("(ll <= a) & (a <= rl)")]',
]

rdict = {}
for i in functions:
    rdict[i] = timeit.timeit(i,setup=setup,number=1000)
    print("%s -> %s s" %(i,rdict[i]))

print("Sorted:")
for w in sorted(rdict, key=rdict.get):
    print(w, rdict[w])

Résultats

Les résultats sont présentés dans le graphique suivant pour un petit tableau (en haut, la solution la plus rapide), comme l'a noté @EZLearner, ils peuvent varier en fonction de la taille du tableau. sorted slice pourrait être plus rapide pour les tableaux plus grands, mais elle nécessite que votre tableau soit trié, pour les tableaux avec plus de 10 M d'entrées ne.evaluate pourrait être une option. Il est donc toujours préférable d'effectuer ce test avec un tableau de la même taille que le vôtre :

Si, au lieu des index, vous voulez extraire les valeurs, vous pouvez effectuer les tests en utilisant les fonctions2 mais les résultats sont presque les mêmes.

Answer 4

J'ai pensé que je devais ajouter ceci parce que le a dans l'exemple que vous avez donné est trié :

import numpy as np
a = [1, 3, 5, 6, 9, 10, 14, 15, 56] 
start = np.searchsorted(a, 6, 'left')
end = np.searchsorted(a, 10, 'right')
rng = np.arange(start, end)
rng
# array([3, 4, 5])

Answer 5

a = np.array([1,2,3,4,5,6,7,8,9])
b = a[(a>2) & (a<8)]