web-dev-qa-db-fra.com

Trouver des lignes correspondantes dans un tableau numpy en 2 dimensions

Je voudrais obtenir l'index d'un tableau Numpy à 2 dimensions qui correspond à une ligne. Par exemple, mon tableau est le suivant:

vals = np.array([[0, 0],
                 [1, 0],
                 [2, 0],
                 [0, 1],
                 [1, 1],
                 [2, 1],
                 [0, 2],
                 [1, 2],
                 [2, 2],
                 [0, 3],
                 [1, 3],
                 [2, 3],
                 [0, 0],
                 [1, 0],
                 [2, 0],
                 [0, 1],
                 [1, 1],
                 [2, 1],
                 [0, 2],
                 [1, 2],
                 [2, 2],
                 [0, 3],
                 [1, 3],
                 [2, 3]])

Je voudrais obtenir l'index qui correspond à la ligne [0, 1] qui est l'index 3 et 15. Quand je fais quelque chose comme numpy.where(vals == [0 ,1]) j'obtiens ...

(array([ 0,  3,  3,  4,  5,  6,  9, 12, 15, 15, 16, 17, 18, 21]), array([0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0]))

Je veux un tableau d'index ([3, 15]).

pythonnumpyscipy
31
b10hazard

Vous avez besoin de la fonction np.where pour obtenir les index:

>>> np.where((vals == (0, 1)).all(axis=1))
(array([ 3, 15]),)

Ou, comme l'indique la documentation:

Si seule la condition est donnée, retournez condition.nonzero()

Vous pouvez appeler directement .nonzero() sur le tableau retourné par .all:

>>> (vals == (0, 1)).all(axis=1).nonzero()
(array([ 3, 15]),)

Pour démonter cela:

>>> vals == (0, 1)
array([[ True, False],
       [False, False],
       ...
       [ True, False],
       [False, False],
       [False, False]], dtype=bool)

et appeler la méthode .all sur ce tableau (avec axis=1) vous donne True où les deux sont True:

>>> (vals == (0, 1)).all(axis=1)
array([False, False, False,  True, False, False, False, False, False,
       False, False, False, False, False, False,  True, False, False,
       False, False, False, False, False, False], dtype=bool)

et pour obtenir quels index sont True:

>>> np.where((vals == (0, 1)).all(axis=1))
(array([ 3, 15]),)

ou

>>> (vals == (0, 1)).all(axis=1).nonzero()
(array([ 3, 15]),)

Je trouve ma solution un peu plus lisible, mais comme le souligne unutbu, ce qui suit peut être plus rapide et renvoie la même valeur que (vals == (0, 1)).all(axis=1):

>>> (vals[:, 0] == 0) & (vals[:, 1] == 1)
40
Aaron Hall
In [5]: np.where((vals[:,0] == 0) & (vals[:,1]==1))[0]
Out[5]: array([ 3, 15])

Je ne sais pas pourquoi, mais c'est beaucoup plus rapide que
np.where((vals == (0, 1)).all(axis=1)):

In [34]: vals2 = np.tile(vals, (1000,1))

In [35]: %timeit np.where((vals2 == (0, 1)).all(axis=1))[0]
1000 loops, best of 3: 808 µs per loop

In [36]: %timeit np.where((vals2[:,0] == 0) & (vals2[:,1]==1))[0]
10000 loops, best of 3: 152 µs per loop
5
unutbu

En utilisant le package numpy_indexed , vous pouvez simplement écrire:

import numpy_indexed as npi
print(np.flatnonzero(npi.contains([[0, 1]], vals)))
3
Eelco Hoogendoorn