web-dev-qa-db-fra.com

Sélection de lignes dans un ndarray NumPy

Je souhaite sélectionner uniquement certaines lignes d'un tableau NumPy en fonction de la valeur dans 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 ne voulais que 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 lorsqu'il y a plus d'une valeur souhaitée?

La liste recherchée peut être de longueur arbitraire. Par exemple, je peux vouloir toutes les lignes où la deuxième colonne est soit 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 reconvertir en un tableau et semble trop compliqué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]])

Existe-t-il une meilleure façon de le faire dans NumPy lui-même qui me manque?

pythonnumpy
24
Raja Selvaraj
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.

16
Amnon

La solution suivante devrait être plus rapide que la solution d'Amnon car wanted s'agrandit:

# 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 uniquement ne fois (au lieu de autant de len(wanted) fois que dans la réponse d'Amnon). Il utilise également la recherche rapide d'éléments intégrés 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 qu'un élément wanted correspond, les éléments restants n'ont pas à être testés (contrairement à l'approche "logique ou" d'Amnon, étaient tous les éléments dans wanted sont testés quoi qu'il arrive).

Alternativement, vous pouvez utiliser le one-liner suivant, qui ne passe également par votre tableau qu'une seule fois:

test[numpy.apply_along_axis(lambda x: x[1] in wanted, 1, test)]

C'est beaucoup plus lent, cependant, car cela extrait l'élément dans la deuxième colonne à chaque itération (au lieu de le faire en une seule fois, comme dans la première solution de cette réponse).

31
Eric O Lebigot

numpy.in1d est ce que vous recherchez:

print test[numpy.in1d(test[:,1], wanted)]

Elle devrait facilement être la solution la plus rapide si elle est importante; en plus, c'est le plus lisible, dis-je.

10
Eelco Hoogendoorn

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), :]
0
Antony Hatchkins