Numpy: trouver rapidement le premier index de valeur

Comment puis-je trouver l'index de la première occurrence d'un nombre dans un tableau Numpy? La vitesse est importante pour moi. Je ne suis pas intéressé par les réponses suivantes car elles analysent l'ensemble du tableau et ne s'arrêtent pas quand elles trouvent la première occurrence:

itemindex = numpy.where(array==item)[0][0]
nonzero(array == item)[0][0]

Note 1: aucune des réponses à cette question ne semble pertinente. Existe-t-il une fonction Numpy pour renvoyer le premier index de quelque chose dans un tableau?

Note 2: l'utilisation d'une méthode compilée en C est préférable à une boucle Python.

Il y a une demande de fonctionnalité pour cela prévue pour Numpy 2.0.0: https://github.com/numpy/numpy/issues/2269

Bien qu'il soit bien trop tard pour vous, mais pour référence future: utiliser numba ( 1 ) est le moyen le plus simple jusqu'à ce que numpy l'implémente. Si vous utilisez la distribution anaconda python, elle devrait déjà être installée. Le code sera compilé donc il sera rapide.

@jit(nopython=True)
def find_first(item, vec):
    """return the index of the first occurence of item in vec"""
    for i in xrange(len(vec)):
        if item == vec[i]:
            return i
    return -1

puis:

>>> a = array([1,7,8,32])
>>> find_first(8,a)
2

J'ai fait une référence pour plusieurs méthodes:

• argwhere
• nonzero comme dans la question
• .tostring() comme dans la réponse de @Rob Reilink
• boucle python
• Boucle de Fortran

Les codes Python et Fortran sont disponibles. J'ai sauté les moins prometteurs comme la conversion en liste.

Les résultats sur une échelle logarithmique. L'axe X est la position de l'aiguille (il faut plus de temps pour savoir si elle est plus loin dans le tableau); la dernière valeur est une aiguille qui n'est pas dans le tableau. L'axe Y est le temps de le trouver.

Le tableau contenait 1 million d'éléments et les tests ont été exécutés 100 fois. Les résultats fluctuent encore un peu, mais la tendance qualitative est claire: Python et f2py s'arrêtent au premier élément, ils évoluent donc différemment. Python devient trop lent si l'aiguille n'est pas dans le premier 1%, alors qu'il f2pyest rapide (mais vous devez le compiler).

Pour résumer, f2py est la solution la plus rapide , surtout si l'aiguille apparaît assez tôt.

Ce n'est pas intégré ce qui est ennuyeux, mais ce n'est vraiment que 2 minutes de travail. Ajoutez ceci à un fichier appelé search.f90:

subroutine find_first(needle, haystack, haystack_length, index)
    implicit none
    integer, intent(in) :: needle
    integer, intent(in) :: haystack_length
    integer, intent(in), dimension(haystack_length) :: haystack
!f2py intent(inplace) haystack
    integer, intent(out) :: index
    integer :: k
    index = -1
    do k = 1, haystack_length
        if (haystack(k)==needle) then
            index = k - 1
            exit
        endif
    enddo
end

Si vous recherchez autre chose que integer, changez simplement le type. Puis compilez en utilisant:

f2py -c -m search search.f90

après quoi vous pouvez faire (à partir de Python):

import search
print(search.find_first.__doc__)
a = search.find_first(your_int_needle, your_int_array)

Vous pouvez convertir un tableau booléen en chaîne Python en utilisant array.tostring()puis en utilisant la méthode find ():

(array==item).tostring().find('\x01')

Cela implique cependant de copier les données, car les chaînes Python doivent être immuables. Un avantage est que vous pouvez également rechercher par exemple un front montant en trouvant\x00\x01

Dans le cas de tableaux triés np.searchsortedfonctionne.

Je pense que vous avez rencontré un problème où une méthode différente et une connaissance a priori du tableau aideraient vraiment. Le genre de chose où vous avez une probabilité X de trouver votre réponse dans les Y premiers pour cent des données. Le fractionnement du problème avec l'espoir d'avoir de la chance, puis de le faire en python avec une compréhension de liste imbriquée ou quelque chose.

Ecrire une fonction C pour faire cette force brute n'est pas non plus trop difficile en utilisant des ctypes .

Le code C que j'ai piraté ensemble (index.c):

long index(long val, long *data, long length){
    long ans, i;
    for(i=0;i<length;i++){
        if (data[i] == val)
            return(i);
    }
    return(-999);
}

et le python:

# to compile (mac)
# gcc -shared index.c -o index.dylib
import ctypes
lib = ctypes.CDLL('index.dylib')
lib.index.restype = ctypes.c_long
lib.index.argtypes = (ctypes.c_long, ctypes.POINTER(ctypes.c_long), ctypes.c_long)

import numpy as np
np.random.seed(8675309)
a = np.random.random_integers(0, 100, 10000)
print lib.index(57, a.ctypes.data_as(ctypes.POINTER(ctypes.c_long)), len(a))

et j'en ai 92.

Enveloppez le python dans une fonction appropriée et le tour est joué.

La version C est beaucoup (~ 20x) plus rapide pour cette graine (attention je ne suis pas bon avec timeit)

import timeit
t = timeit.Timer('np.where(a==57)[0][0]', 'import numpy as np; np.random.seed(1); a = np.random.random_integers(0, 1000000, 10000000)')
t.timeit(100)/100
# 0.09761879920959472
t2 = timeit.Timer('lib.index(57, a.ctypes.data_as(ctypes.POINTER(ctypes.c_long)), len(a))', 'import numpy as np; np.random.seed(1); a = np.random.random_integers(0, 1000000, 10000000); import ctypes; lib = ctypes.CDLL("index.dylib"); lib.index.restype = ctypes.c_long; lib.index.argtypes = (ctypes.c_long, ctypes.POINTER(ctypes.c_long), ctypes.c_long) ')
t2.timeit(100)/100
# 0.005288000106811523

@tal a déjà présenté une numbafonction pour trouver le premier index mais qui ne fonctionne que pour les tableaux 1D. Avec np.ndenumeratevous pouvez également trouver le premier index dans un tableau de dimensions arbitraires:

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
    return None

Exemple de cas:

>>> arr = np.arange(9).reshape(3,3)
>>> index(arr, 3)
(1, 0)

Les délais montrent que ses performances sont similaires à celles de la solution tals :

arr = np.arange(100000)
%timeit index(arr, 5)           # 1000000 loops, best of 3: 1.88 µs per loop
%timeit find_first(5, arr)      # 1000000 loops, best of 3: 1.7 µs per loop

%timeit index(arr, 99999)       # 10000 loops, best of 3: 118 µs per loop
%timeit find_first(99999, arr)  # 10000 loops, best of 3: 96 µs per loop

Si votre liste est triée , vous pouvez effectuer une recherche d'index très rapide avec le package 'bisect'. C'est O (log (n)) au lieu de O (n).

bisect.bisect(a, x)

trouve x dans le tableau a, nettement plus rapide dans le cas trié que n'importe quelle routine C passant par tous les premiers éléments (pour des listes suffisamment longues).

C'est bon à savoir parfois.

Autant que je sache, seuls np.any et np.all sur des tableaux booléens sont court-circuités.

Dans votre cas, numpy doit parcourir le tableau entier deux fois, une fois pour créer la condition booléenne et une seconde fois pour trouver les indices.

Ma recommandation dans ce cas serait d'utiliser cython. Je pense qu'il devrait être facile d'ajuster un exemple pour ce cas, surtout si vous n'avez pas besoin de beaucoup de flexibilité pour différents types et formes.

J'en avais besoin pour mon travail, alors je me suis appris l'interface C de Python et de Numpy et j'ai écrit la mienne. http://pastebin.com/GtcXuLyd Ce n'est que pour les tableaux 1-D, mais fonctionne pour la plupart des types de données (int, float ou strings) et les tests ont montré qu'il est à nouveau environ 20 fois plus rapide que l'approche attendue en Python pur- engourdi.

Ce problème peut être résolu efficacement en numpy pur en traitant le tableau en morceaux:

def find_first(x):
    idx, step = 0, 32
    while idx < x.size:
        nz, = x[idx: idx + step].nonzero()
        if len(nz): # found non-zero, return it
            return nz[0] + idx
        # move to the next chunk, increase step
        idx += step
        step = min(9600, step + step // 2)
    return -1

Le tableau est traité par bloc de taille step. Le stepplus est l'étape, la plus rapide est le traitement de-gamme mis à zéro ( dans le pire des cas). Plus il est petit, plus le traitement du tableau avec une valeur non nulle au début est rapide. L'astuce consiste à commencer par un petit stepet à l'augmenter de façon exponentielle. De plus, il n'est pas nécessaire de l'augmenter au-dessus d'un certain seuil en raison des avantages limités.

J'ai comparé la solution avec la solution pure ndarary.nonzero et numba contre 10 millions de tableaux de flotteurs.

import numpy as np
from numba import jit
from timeit import timeit

def find_first(x):
    idx, step = 0, 32
    while idx < x.size:
        nz, = x[idx: idx + step].nonzero()
        if len(nz):
            return nz[0] + idx
        idx += step
        step = min(9600, step + step // 2)
    return -1

@jit(nopython=True)
def find_first_numba(vec):
    """return the index of the first occurence of item in vec"""
    for i in range(len(vec)):
        if vec[i]:
            return i
    return -1


SIZE = 10_000_000
# First only
x = np.empty(SIZE)

find_first_numba(x[:10])

print('---- FIRST ----')
x[:] = 0
x[0] = 1
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0][0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=1000), 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=1000), 'ms')

print('---- LAST ----')
x[:] = 0
x[-1] = 1
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0][0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=100)*10, 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=100)*10, 'ms')

print('---- NONE ----')
x[:] = 0
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=100)*10, 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=100)*10, 'ms')

print('---- ALL ----')
x[:] = 1
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0][0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=100)*10, 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=100)*10, 'ms')

Et les résultats sur ma machine:

---- FIRST ----
ndarray.nonzero 54.733994480002366 ms
find_first 0.0013148509997336078 ms
find_first_numba 0.0002839310000126716 ms
---- LAST ----
ndarray.nonzero 54.56336712999928 ms
find_first 25.38929685000312 ms
find_first_numba 8.022820680002951 ms
---- NONE ----
ndarray.nonzero 24.13432420999925 ms
find_first 25.345200140000088 ms
find_first_numba 8.154927100003988 ms
---- ALL ----
ndarray.nonzero 55.753537260002304 ms
find_first 0.0014760300018679118 ms
find_first_numba 0.0004358099977253005 ms

Pure ndarray.nonzeroest définitivement plus lâche. La solution numba est environ 5 fois plus rapide dans le meilleur des cas. C'est environ 3 fois plus rapide dans le pire des cas.

Si vous recherchez le premier élément non nul, vous pouvez utiliser un hack suivant:

idx = x.view(bool).argmax() // x.itemsize
idx = idx if x[idx] else -1

C'est une solution "numpy-pure" très rapide mais elle échoue dans certains cas décrits ci-dessous.

La solution tire parti du fait que pratiquement toutes les représentations de zéro pour les types numériques sont constituées d' 0octets. Cela s'applique également à Numpy bool. Dans les versions récentes de numpy, la argmax()fonction utilise la logique de court-circuit lors du traitement du booltype. La taille de boolest de 1 octet.

Il faut donc:

• créer une vue du tableau sous la forme bool. Aucune copie n'est créée
• utiliser argmax()pour trouver le premier octet différent de zéro en utilisant la logique de court-circuit
• recalculer l'offset de cet octet à l'index du premier élément non nul par division entière (opérateur //) de l'offset par une taille d'un seul élément exprimée en octets ( x.itemsize)
• vérifier si x[idx]est réellement non nul pour identifier le cas où aucun non nul n'est présent

J'ai fait une référence par rapport à la solution numba et je l'ai construite np.nonzero.

import numpy as np
from numba import jit
from timeit import timeit

def find_first(x):
    idx = x.view(bool).argmax() // x.itemsize
    return idx if x[idx] else -1

@jit(nopython=True)
def find_first_numba(vec):
    """return the index of the first occurence of item in vec"""
    for i in range(len(vec)):
        if vec[i]:
            return i
    return -1


SIZE = 10_000_000
# First only
x = np.empty(SIZE)

find_first_numba(x[:10])

print('---- FIRST ----')
x[:] = 0
x[0] = 1
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0][0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=1000), 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=1000), 'ms')

print('---- LAST ----')
x[:] = 0
x[-1] = 1
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0][0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=100)*10, 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=100)*10, 'ms')

print('---- NONE ----')
x[:] = 0
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=100)*10, 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=100)*10, 'ms')

print('---- ALL ----')
x[:] = 1
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0][0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=100)*10, 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=100)*10, 'ms')

Le résultat sur ma machine est:

---- FIRST ----
ndarray.nonzero 57.63976670001284 ms
find_first 0.0010841979965334758 ms
find_first_numba 0.0002308919938514009 ms
---- LAST ----
ndarray.nonzero 58.96685277999495 ms
find_first 5.923203580023255 ms
find_first_numba 8.762269750004634 ms
---- NONE ----
ndarray.nonzero 25.13398071998381 ms
find_first 5.924289370013867 ms
find_first_numba 8.810063839919167 ms
---- ALL ----
ndarray.nonzero 55.181210660084616 ms
find_first 0.001246920000994578 ms
find_first_numba 0.00028766007744707167 ms

La solution est 33% plus rapide que numba et elle est "numpy-pure".

Les désavantages:

• ne fonctionne pas pour les types numpy acceptables comme object
• échoue pour un zéro négatif qui apparaît parfois dans floatou des doublecalculs

En tant qu'utilisateur de longue date de matlab, je cherche depuis un certain temps une solution efficace à ce problème. Enfin, motivé par des discussions et des propositions dans ce fil, j'ai essayé de trouver une solution qui implémente une API similaire à ce qui a été suggéré ici , ne supportant pour le moment que les tableaux 1D.

Vous l'utiliseriez comme ça

import numpy as np
import utils_find_1st as utf1st
array = np.arange(100000)
item = 1000
ind = utf1st.find_1st(array, item, utf1st.cmp_larger_eq)

Les opérateurs de condition pris en charge sont: cmp_equal, cmp_not_equal, cmp_larger, cmp_smaller, cmp_larger_eq, cmp_smaller_eq. Pour plus d'efficacité, l'extension est écrite en c.

Vous trouvez la source, les benchmarks et d'autres détails ici:

https://pypi.python.org/pypi?name=py_find_1st&:action=display

pour l'utilisation dans notre équipe (anaconda sur linux et macos) j'ai fait un installeur anaconda qui simplifie l'installation, vous pouvez l'utiliser comme décrit ici

https://anaconda.org/roebel/py_find_1st

Sachez simplement que si vous effectuez une séquence de recherches, le gain de performances en faisant quelque chose d'intelligent comme la conversion en chaîne peut être perdu dans la boucle externe si la dimension de recherche n'est pas assez grande. Voyez comment les performances de l'itération de find1 qui utilise l'astuce de conversion de chaîne proposée ci-dessus et find2 qui utilise argmax le long de l'axe interne (plus un ajustement pour garantir qu'une non-correspondance renvoie -1)

import numpy,time
def find1(arr,value):
    return (arr==value).tostring().find('\x01')

def find2(arr,value): #find value over inner most axis, and return array of indices to the match
    b = arr==value
    return b.argmax(axis=-1) - ~(b.any())


for size in [(1,100000000),(10000,10000),(1000000,100),(10000000,10)]:
    print(size)
    values = numpy.random.choice([0,0,0,0,0,0,0,1],size=size)
    v = values>0

    t=time.time()
    numpy.apply_along_axis(find1,-1,v,1)
    print('find1',time.time()-t)

    t=time.time()
    find2(v,1)
    print('find2',time.time()-t)

les sorties

(1, 100000000)
('find1', 0.25300002098083496)
('find2', 0.2780001163482666)
(10000, 10000)
('find1', 0.46200013160705566)
('find2', 0.27300000190734863)
(1000000, 100)
('find1', 20.98099994659424)
('find2', 0.3040001392364502)
(10000000, 10)
('find1', 206.7590000629425)
('find2', 0.4830000400543213)

Cela dit, une recherche écrite en C serait au moins un peu plus rapide que l'une ou l'autre de ces approches

que dis-tu de ça

import numpy as np
np.amin(np.where(array==item))

Vous pouvez convertir votre tableau en un listet utiliser sa index()méthode:

i = list(array).index(item)

Pour autant que je sache, il s'agit d'une méthode compilée en C.