Numpy: trova velocemente il primo indice di valore

Come posso trovare l'indice della prima occorrenza di un numero in un array Numpy? La velocità è importante per me. Non sono interessato alle seguenti risposte perché scansionano l'intero array e non si fermano quando trovano la prima occorrenza:

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

Nota 1: nessuna delle risposte da questa domanda sembra rilevante Esiste una funzione Numpy per restituire il primo indice di qualcosa in un array?

Nota 2: l'utilizzo di un metodo compilato in C è preferibile a un ciclo Python.

C'è una richiesta di funzionalità per questo programmata per Numpy 2.0.0: https://github.com/numpy/numpy/issues/2269

Anche se è troppo tardi per te, ma per riferimento futuro: usare numba ( 1 ) è il modo più semplice finché numpy non lo implementa. Se usi la distribuzione anaconda python dovrebbe essere già installata. Il codice verrà compilato in modo che sia veloce.

@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

e poi:

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

Ho creato un punto di riferimento per diversi metodi:

• argwhere
• nonzero come nella domanda
• .tostring() come nella risposta di @Rob Reilink
• ciclo Python
• Fortran loop

Il Python e Fortran codice sono disponibili. Ho saltato quelli poco promettenti come la conversione in un elenco.

I risultati su scala logaritmica. L'asse X è la posizione dell'ago (ci vuole più tempo per trovare se è più in basso nella matrice); l'ultimo valore è un ago che non è nell'array. L'asse Y è il momento di trovarlo.

L'array aveva 1 milione di elementi e i test sono stati eseguiti 100 volte. I risultati oscillano ancora un po ', ma la tendenza qualitativa è chiara: Python e f2py si chiudono al primo elemento in modo da scalare in modo diverso. Python diventa troppo lento se l'ago non è nel primo 1%, mentre f2pyè veloce (ma è necessario compilarlo).

Per riassumere, f2py è la soluzione più veloce , soprattutto se l'ago appare abbastanza presto.

Non è integrato, il che è fastidioso, ma in realtà sono solo 2 minuti di lavoro. Aggiungere questo in un file chiamato 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

Se stai cercando qualcosa di diverso integer, cambia semplicemente il tipo. Quindi compila usando:

f2py -c -m search search.f90

dopodiché puoi fare (da Python):

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

Puoi convertire un array booleano in una stringa Python usando array.tostring()e poi usando il metodo find ():

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

Ciò comporta la copia dei dati, tuttavia, poiché le stringhe Python devono essere immutabili. Un vantaggio è che puoi anche cercare, ad esempio, un fronte di salita trovando\x00\x01

In caso di array ordinati np.searchsortedfunziona.

Penso che tu abbia riscontrato un problema in cui un metodo diverso e una certa conoscenza a priori dell'array sarebbero davvero utili. Il tipo di cosa in cui hai una probabilità X di trovare la tua risposta nel primo Y percento dei dati. La divisione del problema con la speranza di essere fortunati, quindi farlo in Python con una comprensione di liste annidate o qualcosa del genere.

Anche scrivere una funzione C per eseguire questa forza bruta non è troppo difficile usando ctypes .

Il codice C che ho hackerato insieme (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);
}

e il pitone:

# 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))

e ottengo 92.

Avvolgi il pitone in una funzione appropriata e il gioco è fatto.

La versione C è molto (~ 20 volte) più veloce per questo seme (avvertimento che non sono bravo con il tempo)

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 ha già presentato una numbafunzione per trovare il primo indice, ma funziona solo per array 1D. Con np.ndenumeratepuoi anche trovare il primo indice in un array dimensionale arbitrario:

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

Caso di esempio:

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

I tempi mostrano che è simile nelle prestazioni alla soluzione 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

Se la tua lista è ordinata , puoi ottenere una ricerca dell'indice molto rapida con il pacchetto 'bisect'. È O (log (n)) invece di O (n).

bisect.bisect(a, x)

trova x nell'array a, decisamente più veloce nel caso ordinato di qualsiasi routine C che passa attraverso tutti i primi elementi (per elenchi abbastanza lunghi).

È bello sapere a volte.

Per quanto ne so, solo np.any e np.all sugli array booleani sono in cortocircuito.

Nel tuo caso, numpy deve passare attraverso l'intero array due volte, una volta per creare la condizione booleana e una seconda volta per trovare gli indici.

La mia raccomandazione in questo caso sarebbe quella di utilizzare cython. Penso che dovrebbe essere facile modificare un esempio per questo caso, soprattutto se non è necessaria molta flessibilità per diversi tipi e forme.

Ne avevo bisogno per il mio lavoro, quindi ho imparato da solo l'interfaccia C di Python e Numpy e ho scritto la mia. http://pastebin.com/GtcXuLyd È solo per array 1-D, ma funziona per la maggior parte dei tipi di dati (int, float o stringhe) e il test ha dimostrato che è di nuovo circa 20 volte più veloce dell'approccio previsto in puro Python- NumPy.

Questo problema può essere risolto efficacemente in puro numpy elaborando l'array in blocchi:

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

L'array viene elaborato in blocchi di dimensione step. Più steplungo è il passo, più veloce sarà l'elaborazione di un array a zero (caso peggiore). Più è piccolo, più veloce è l'elaborazione di un array con un valore diverso da zero all'inizio. Il trucco è iniziare con un piccolo stepe aumentarlo in modo esponenziale. Inoltre, non è necessario incrementarlo oltre una certa soglia a causa dei vantaggi limitati.

Ho confrontato la soluzione con la soluzione ndarary.nonzero e numba pura con 10 milioni di array di float.

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')

E risultati sulla mia macchina:

---- 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.nonzeroè decisamente più sciolto. La soluzione numba è circa 5 volte più veloce per il caso migliore. È circa 3 volte più veloce nel peggiore dei casi.

Se stai cercando il primo elemento diverso da zero puoi usare un seguente trucco:

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

È una soluzione "numpy-pure" molto veloce ma fallisce per alcuni casi discussi di seguito.

La soluzione trae vantaggio dal fatto che praticamente tutta la rappresentazione di zero per i tipi numerici è costituita da 0byte. Si applica anche a numpy bool. Nelle versioni recenti di numpy, la argmax()funzione utilizza la logica di cortocircuito durante l'elaborazione del booltipo. La dimensione di boolè 1 byte.

Quindi è necessario:

• creare una vista dell'array come bool. Nessuna copia viene creata
• utilizzare argmax()per trovare il primo byte diverso da zero utilizzando la logica di cortocircuito
• ricalcola l'offset di questo byte all'indice del primo elemento diverso da zero per divisione intera (operatore //) dell'offset per una dimensione di un singolo elemento espressa in byte ( x.itemsize)
• controlla se x[idx]è effettivamente diverso da zero per identificare il caso in cui non è presente alcun diverso da zero

Ho fatto alcuni benchmark contro la soluzione numba e l'ho costruito 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')

I risultati sulla mia macchina sono:

---- 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 soluzione è il 33% più veloce di numba ed è "numpy-pure".

Gli svantaggi:

• non funziona per tipi accettabili come numpy object
• fallisce per zero negativo che appare occasionalmente in floato doublecalcoli

Come utente di matlab di lunga data, ho cercato una soluzione efficiente a questo problema per un po 'di tempo. Infine, motivato da discussioni e proposizioni in questo thread, ho cercato di trovare una soluzione che implementasse un'API simile a quanto qui suggerito , supportando per il momento solo array 1D.

Lo useresti in questo modo

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)

Gli operatori di condizione supportati sono: cmp_equal, cmp_not_equal, cmp_larger, cmp_smaller, cmp_larger_eq, cmp_smaller_eq. Per efficienza l'estensione è scritta in c.

Puoi trovare la fonte, i benchmark e altri dettagli qui:

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

per l'utilizzo nel nostro team (anaconda su linux e macos) ho realizzato un installer anaconda che semplifica l'installazione, puoi usarlo come descritto qui

https://anaconda.org/roebel/py_find_1st

Solo una nota che se stai eseguendo una sequenza di ricerche, il guadagno di prestazioni derivante da qualcosa di intelligente come la conversione in stringa, potrebbe essere perso nel ciclo esterno se la dimensione di ricerca non è abbastanza grande. Guarda come le prestazioni dell'iterazione di find1 che utilizza il trucco di conversione delle stringhe proposto sopra e find2 che utilizza argmax lungo l'asse interno (più un aggiustamento per garantire che una non corrispondenza restituisca -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)

uscite

(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)

Detto questo, una ricerca scritta in C sarebbe almeno un po 'più veloce di uno di questi approcci

cosa ne pensi di questo

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

Puoi convertire il tuo array in un liste utilizzare il suo index()metodo:

i = list(array).index(item)

Per quanto ne so, questo è un metodo compilato in C.