Existe-t-il une fonction NumPy pour retourner le premier index de quelque chose dans un tableau?

Je sais qu'il existe une méthode pour qu'une liste Python retourne le premier index de quelque chose:

>>> l = [1, 2, 3]
>>> l.index(2)
1

Y a-t-il quelque chose comme ça pour les tableaux NumPy?

Oui, voici la réponse donnée un tableau NumPy, array , et une valeur item, pour rechercher:

itemindex = numpy.where(array==item)

Le résultat est un tuple avec d'abord tous les indices de ligne, puis tous les indices de colonne.

Par exemple, si un tableau a deux dimensions et qu'il contenait votre article à deux endroits,

array[itemindex[0][0]][itemindex[1][0]]

serait égal à votre article et il en serait de même

array[itemindex[0][1]][itemindex[1][1]]

numpy.where

Si vous avez besoin de l'index de la première occurrence d' une seule valeur , vous pouvez utiliser nonzero(ou where, ce qui revient au même dans ce cas):

>>> t = array([1, 1, 1, 2, 2, 3, 8, 3, 8, 8])
>>> nonzero(t == 8)
(array([6, 8, 9]),)
>>> nonzero(t == 8)[0][0]
6

Si vous avez besoin du premier index de chacune des nombreuses valeurs , vous pouvez évidemment faire la même chose que ci-dessus à plusieurs reprises, mais il existe une astuce qui peut être plus rapide. Ce qui suit trouve les indices du premier élément de chaque sous- séquence :

>>> nonzero(r_[1, diff(t)[:-1]])
(array([0, 3, 5, 6, 7, 8]),)

Notez qu'il trouve le début de la sous-séquence de 3 et des deux sous-séquences de 8:

[ 1 , 1, 1, 2 , 2, 3 , 8 , 3 , 8 , 8]

C'est donc légèrement différent de trouver la première occurrence de chaque valeur. Dans votre programme, vous pourrez peut-être travailler avec une version triée de tpour obtenir ce que vous voulez:

>>> st = sorted(t)
>>> nonzero(r_[1, diff(st)[:-1]])
(array([0, 3, 5, 7]),)

Vous pouvez également convertir un tableau NumPy en liste dans l'air et obtenir son index. Par exemple,

l = [1,2,3,4,5] # Python list
a = numpy.array(l) # NumPy array
i = a.tolist().index(2) # i will return index of 2
print i

Il imprimera 1.

Juste pour ajouter un très performant et pratique numbaalternative basée sur np.ndenumeratepour trouver le premier index:

from numba import njit
import numpy as np

@njit
def index(array, item):
    for idx, val in np.ndenumerate(array):
        if val == item:
            return idx
    # If no item was found return None, other return types might be a problem due to
    # numbas type inference.

C'est assez rapide et traite naturellement des tableaux multidimensionnels :

>>> arr1 = np.ones((100, 100, 100))
>>> arr1[2, 2, 2] = 2

>>> index(arr1, 2)
(2, 2, 2)

>>> arr2 = np.ones(20)
>>> arr2[5] = 2

>>> index(arr2, 2)
(5,)

Cela peut être beaucoup plus rapide (car il court-circuite l'opération) que toute approche utilisant np.whereou np.nonzero.

Cependant, cela np.argwherepourrait également fonctionner avec élégance avec les tableaux multidimensionnels (vous auriez besoin de le convertir manuellement en un tuple et il n'est pas court-circuité), mais il échouerait si aucune correspondance n'était trouvée:

>>> tuple(np.argwhere(arr1 == 2)[0])
(2, 2, 2)
>>> tuple(np.argwhere(arr2 == 2)[0])
(5,)

Si vous allez utiliser ceci comme index dans quelque chose d'autre, vous pouvez utiliser des index booléens si les tableaux sont diffusables; vous n'avez pas besoin d'indices explicites. La façon la plus simple de procéder consiste à simplement indexer en fonction d'une valeur de vérité.

other_array[first_array == item]

Toute opération booléenne fonctionne:

a = numpy.arange(100)
other_array[first_array > 50]

La méthode non nulle prend aussi des booléens:

index = numpy.nonzero(first_array == item)[0][0]

Les deux zéros correspondent au tuple d'indices (en supposant que first_array est 1D), puis au premier élément du tableau d'indices.

l.index(x)renvoie le plus petit i tel que i soit l'indice de la première occurrence de x dans la liste.

On peut supposer en toute sécurité que la index()fonction en Python est implémentée de manière à ce qu'elle s'arrête après avoir trouvé la première correspondance, et cela se traduit par une performance moyenne optimale.

Pour trouver un élément s'arrêtant après la première correspondance dans un tableau NumPy, utilisez un itérateur ( ndenumerate ).

In [67]: l=range(100)

In [68]: l.index(2)
Out[68]: 2

Tableau NumPy:

In [69]: a = np.arange(100)

In [70]: next((idx for idx, val in np.ndenumerate(a) if val==2))
Out[70]: (2L,)

Notez que les deux méthodes index()et nextrenvoient une erreur si l'élément est introuvable. Avec next, on peut utiliser un deuxième argument pour retourner une valeur spéciale au cas où l'élément ne serait pas trouvé, par exemple

In [77]: next((idx for idx, val in np.ndenumerate(a) if val==400),None)

Il existe d' autres fonctions dans NumPy ( argmax, whereet nonzero) qui peut être utilisé pour trouver un élément dans un tableau, mais ils ont tous l'inconvénient de passer par le tableau entier à la recherche de toutes les occurrences, donc pas optimisé pour trouver le premier élément. Notez également cela whereet nonzerorenvoyez des tableaux, vous devez donc sélectionner le premier élément pour obtenir l'index.

In [71]: np.argmax(a==2)
Out[71]: 2

In [72]: np.where(a==2)
Out[72]: (array([2], dtype=int64),)

In [73]: np.nonzero(a==2)
Out[73]: (array([2], dtype=int64),)

Comparaison de temps

Vérifier simplement que pour les grands tableaux, la solution utilisant un itérateur est plus rapide lorsque l'élément recherché est au début du tableau (en utilisant %timeitdans le shell IPython):

In [285]: a = np.arange(100000)

In [286]: %timeit next((idx for idx, val in np.ndenumerate(a) if val==0))
100000 loops, best of 3: 17.6 µs per loop

In [287]: %timeit np.argmax(a==0)
1000 loops, best of 3: 254 µs per loop

In [288]: %timeit np.where(a==0)[0][0]
1000 loops, best of 3: 314 µs per loop

Il s'agit d'un problème ouvert avec NumPy GitHub .

Voir aussi: Numpy: trouver rapidement le premier indice de valeur

Pour les tableaux triés unidimensionnels , il serait beaucoup plus simple et efficace O (log (n)) d'utiliser numpy.searchsorted qui renvoie un entier NumPy (position). Par exemple,

arr = np.array([1, 1, 1, 2, 3, 3, 4])
i = np.searchsorted(arr, 3)

Assurez-vous simplement que le tableau est déjà trié

Vérifiez également si l'index retourné i contient réellement l'élément recherché, car l'objectif principal de searchsorted est de trouver des indices dans lesquels les éléments doivent être insérés pour maintenir l'ordre.

if arr[i] == 3:
    print("present")
else:
    print("not present")

Pour indexer sur n'importe quel critère, vous pouvez donc quelque chose comme ceci:

In [1]: from numpy import *
In [2]: x = arange(125).reshape((5,5,5))
In [3]: y = indices(x.shape)
In [4]: locs = y[:,x >= 120] # put whatever you want in place of x >= 120
In [5]: pts = hsplit(locs, len(locs[0]))
In [6]: for pt in pts:
   .....:         print(', '.join(str(p[0]) for p in pt))
4, 4, 0
4, 4, 1
4, 4, 2
4, 4, 3
4, 4, 4

Et voici une fonction rapide pour faire ce que fait list.index (), sauf qu'il ne déclenche pas d'exception s'il n'est pas trouvé. Attention - cela est probablement très lent sur les grands tableaux. Vous pouvez probablement corriger cela sur des tableaux si vous préférez l'utiliser comme méthode.

def ndindex(ndarray, item):
    if len(ndarray.shape) == 1:
        try:
            return [ndarray.tolist().index(item)]
        except:
            pass
    else:
        for i, subarray in enumerate(ndarray):
            try:
                return [i] + ndindex(subarray, item)
            except:
                pass

In [1]: ndindex(x, 103)
Out[1]: [4, 0, 3]

Pour les tableaux 1D, je recommanderais np.flatnonzero(array == value)[0], ce qui équivaut aux deux np.nonzero(array == value)[0][0]et np.where(array == value)[0][0]évite la laideur de déballer un tuple à 1 élément.

Une alternative à la sélection du premier élément dans np.where () consiste à utiliser une expression de générateur avec énumération, telle que:

>>> import numpy as np
>>> x = np.arange(100)   # x = array([0, 1, 2, 3, ... 99])
>>> next(i for i, x_i in enumerate(x) if x_i == 2)
2

Pour un tableau à deux dimensions, on ferait:

>>> x = np.arange(100).reshape(10,10)   # x = array([[0, 1, 2,... 9], [10,..19],])
>>> next((i,j) for i, x_i in enumerate(x) 
...            for j, x_ij in enumerate(x_i) if x_ij == 2)
(0, 2)

L'avantage de cette approche est qu'elle arrête de vérifier les éléments du tableau une fois la première correspondance trouvée, tandis que np.where vérifie la correspondance de tous les éléments. Une expression de générateur serait plus rapide s'il y a correspondance au début du tableau.

Il existe de nombreuses opérations dans NumPy qui pourraient peut-être être regroupées pour y parvenir. Cela renverra des indices d'éléments égaux à l'élément:

numpy.nonzero(array - item)

Vous pouvez ensuite prendre les premiers éléments des listes pour obtenir un seul élément.

Le paquet numpy_indexed (avertissement, je suis son auteur) contient un équivalent vectorisé de list.index pour numpy.ndarray; C'est:

sequence_of_arrays = [[0, 1], [1, 2], [-5, 0]]
arrays_to_query = [[-5, 0], [1, 0]]

import numpy_indexed as npi
idx = npi.indices(sequence_of_arrays, arrays_to_query, missing=-1)
print(idx)   # [2, -1]

Cette solution a des performances vectorisées, se généralise en ndarrays et a différentes manières de traiter les valeurs manquantes.

Remarque: c'est pour la version python 2.7

Vous pouvez utiliser une fonction lambda pour résoudre le problème, et elle fonctionne à la fois sur le tableau et la liste NumPy.

your_list = [11, 22, 23, 44, 55]
result = filter(lambda x:your_list[x]>30, range(len(your_list)))
#result: [3, 4]

import numpy as np
your_numpy_array = np.array([11, 22, 23, 44, 55])
result = filter(lambda x:your_numpy_array [x]>30, range(len(your_list)))
#result: [3, 4]

Et vous pouvez utiliser

result[0]

pour obtenir le premier index des éléments filtrés.

Pour python 3.6, utilisez

list(result)

au lieu de

result