Sélection de lignes à partir D'un ndarray NumPy
Je veux sélectionner uniquement certaines lignes d'un tableau NumPy basé sur la valeur de la deuxième colonne. Par exemple, ce tableau de test a des entiers de 1 à 10 dans la deuxième colonne.
>>> test = numpy.array([numpy.arange(100), numpy.random.randint(1, 11, 100)]).transpose()
>>> test[:10, :]
array([[ 0, 6],
[ 1, 7],
[ 2, 10],
[ 3, 4],
[ 4, 1],
[ 5, 10],
[ 6, 6],
[ 7, 4],
[ 8, 6],
[ 9, 7]])
Si je voulais seulement des lignes où la deuxième valeur est 4, c'est facile:
>>> test[test[:, 1] == 4]
array([[ 3, 4],
[ 7, 4],
[16, 4],
...
[81, 4],
[83, 4],
[88, 4]])
Mais comment puis-je obtenir le même résultat quand il y a plus d'une valeur désirée?
La liste désirée peut être de longueur arbitraire. Par exemple, je peux vouloir toutes les lignes où la deuxième colonne est 2, 4 ou 6:
>>> wanted = [2, 4, 6]
La seule façon que j'ai trouvée est d'utiliser la compréhension de liste, puis de la convertir en un tableau et semble trop alambiquée, bien que cela fonctionne:
>>> test[numpy.array([test[x, 1] in wanted for x in range(len(test))])]
array([[ 0, 6],
[ 3, 4],
[ 6, 6],
...
[90, 2],
[91, 6],
[92, 2]])
Y a-t-il une meilleure façon de le faire en NumPy lui-même que je manque?
4 réponses
test[numpy.logical_or.reduce([test[:,1] == x for x in wanted])]
Le résultat devrait être plus rapide que la version originale puisque NumPy fait les boucles internes au lieu de Python.
La solution suivante devrait être plus rapide que la solution D'Amnon car wanted devient plus grande:
# Much faster look up than with lists, for larger lists:
wanted_set = set(wanted)
@numpy.vectorize
def selected(elmt): return elmt in wanted_set
# Or: selected = numpy.vectorize(wanted_set.__contains__)
print test[selected(test[:, 1])]
En fait, il a l'avantage de rechercher dans le tableau test seulement Une fois (au lieu d'autant de len(wanted) fois que dans la réponse D'Amnon). Il utilise également l'élément rapide intégré de Python dans sets , qui sont beaucoup plus rapides pour cela que les listes. Il est également rapide car il utilise les boucles rapides de Numpy. Vous obtenez également l'optimisation de l'opérateur in: Une fois un wanted les éléments restants ne doivent pas être testés (par opposition à L'approche "logique ou" D'Amnon, où tous les éléments de wanted sont testés quoi qu'il arrive).
Alternativement, vous pouvez utiliser le One-liner suivant, qui ne traverse également votre tableau qu'une seule fois:
test[numpy.apply_along_axis(lambda x: x[1] in wanted, 1, test)]
Ceci est beaucoup plus lent, car cela extrait l'élément dans la deuxième colonne à chaque itération (au lieu de le faire en une seule passe, comme dans la première solution de cette réponse).
Numpy. in1d est ce que vous recherchez:
print test[numpy.in1d(test[:,1], wanted)]
Il devrait facilement être la solution la plus rapide si voulu est grand; de plus, il est le plus lisible, ID dire.
C'est deux fois plus rapide que la variante D'Amnon pour len (test)=1000:
wanted = (2,4,6)
wanted2 = numpy.expand_dims(wanted, 1)
print test[numpy.any(test[:, 1] == wanted2, 0), :]