Numpy: encuentra el primer índice de valor rápidamente

¿Cómo puedo encontrar el índice de la primera aparición de un número en una matriz Numpy? La velocidad es importante para mi. No me interesan las siguientes respuestas porque escanean toda la matriz y no se detienen cuando encuentran la primera aparición:

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

Nota 1: ninguna de las respuestas de esa pregunta parece relevante ¿Existe una función Numpy para devolver el primer índice de algo en una matriz?

Nota 2: se prefiere utilizar un método compilado en C a un bucle de Python.

Hay una solicitud de función para esto programada para Numpy 2.0.0: https://github.com/numpy/numpy/issues/2269

Aunque es demasiado tarde para usted, pero para referencia futura: usar numba ( 1 ) es la forma más fácil hasta que numpy lo implemente. Si usa la distribución anaconda python, ya debería estar instalada. El código se compilará para que sea rápido.

@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

y entonces:

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

Hice un punto de referencia para varios métodos:

• argwhere
• nonzero como en la pregunta
• .tostring() como en la respuesta de @Rob Reilink
• bucle de Python
• Bucle de Fortran

El código Python y Fortran están disponibles. Me salté los poco prometedores, como convertirlos en una lista.

Los resultados a escala logarítmica. El eje X es la posición de la aguja (se tarda más en encontrar si está más abajo en la matriz); El último valor es una aguja que no está en la matriz. El eje Y es el momento de encontrarlo.

La matriz tenía 1 millón de elementos y las pruebas se ejecutaron 100 veces. Los resultados aún fluctúan un poco, pero la tendencia cualitativa es clara: Python y f2py abandonan el primer elemento, por lo que escalan de manera diferente. Python se vuelve demasiado lento si la aguja no está en el primer 1%, mientras que f2pyes rápido (pero debe compilarlo).

En resumen, f2py es la solución más rápida , especialmente si la aguja aparece bastante pronto.

No está integrado, lo que es molesto, pero en realidad son solo 2 minutos de trabajo. Agregue esto a un archivo llamado 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 está buscando algo diferente a integer, simplemente cambie el tipo. Luego compila usando:

f2py -c -m search search.f90

después de lo cual puede hacer (desde Python):

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

Puede convertir una matriz booleana en una cadena de Python usando array.tostring()y luego usando el método find ():

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

Sin embargo, esto implica copiar los datos, ya que las cadenas de Python deben ser inmutables. Una ventaja es que también puede buscar, por ejemplo, un flanco ascendente al encontrar\x00\x01

En caso de arreglos ordenados, np.searchsortedfunciona.

Creo que ha encontrado un problema en el que un método diferente y algo a priori conocimiento a de la matriz realmente ayudarían. El tipo de cosas en las que tienes una probabilidad X de encontrar tu respuesta en el primer Y por ciento de los datos. La división del problema con la esperanza de tener suerte y luego hacer esto en python con una lista de comprensión anidada o algo así.

Escribir una función C para hacer esta fuerza bruta tampoco es demasiado difícil con ctypes .

El código C que pirateé juntos (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);
}

y la pitón:

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

y obtengo 92.

Envuelva la pitón en una función adecuada y listo.

La versión C es mucho (~ 20x) más rápida para esta semilla (advirtiendo que no soy bueno con 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 ya presentó una numbafunción para encontrar el primer índice, pero eso solo funciona para matrices 1D. Con np.ndenumeratetambién puede encontrar el primer índice en una matriz dimensional arbitraria:

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 de muestra:

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

Los tiempos muestran que es similar en rendimiento a la solución 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 su lista está ordenada , puede lograr una búsqueda de índice muy rápida con el paquete 'bisect'. Es O (log (n)) en lugar de O (n).

bisect.bisect(a, x)

encuentra x en la matriz a, definitivamente más rápido en el caso ordenado que cualquier rutina C que pasa por todos los primeros elementos (para listas lo suficientemente largas).

A veces es bueno saberlo.

Hasta donde yo sé, solo np.any y np.all en matrices booleanas están en cortocircuito.

En su caso, numpy tiene que pasar por toda la matriz dos veces, una para crear la condición booleana y una segunda para encontrar los índices.

Mi recomendación en este caso sería utilizar cython. Creo que debería ser fácil ajustar un ejemplo para este caso, especialmente si no necesita mucha flexibilidad para diferentes tipos y formas.

Necesitaba esto para mi trabajo, así que me enseñé a mí mismo la interfaz C de Python y Numpy y escribí la mía propia. http://pastebin.com/GtcXuLyd Es solo para matrices 1-D, pero funciona para la mayoría de los tipos de datos (int, float o strings) y las pruebas han demostrado que es nuevamente unas 20 veces más rápido que el enfoque esperado en Python puro. numpy.

Este problema se puede resolver de manera efectiva en números puros procesando la matriz en fragmentos:

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

La matriz se procesa en trozos de tamaño step. Cuanto stepmás largo sea el paso, más rápido será el procesamiento de la matriz con cero (el peor de los casos). Cuanto más pequeño es, más rápido se procesa la matriz con un valor distinto de cero al principio. El truco consiste en empezar con una pequeña cantidad stepy aumentarla exponencialmente. Además, no es necesario incrementarlo por encima de algún umbral debido a los beneficios limitados.

He comparado la solución con la solución ndarary.nonzero y numba pura con 10 millones de arreglos flotantes.

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

Y resultados en mi máquina:

---- 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.nonzeroes definitivamente más suelto. La solución numba es alrededor de 5 veces más rápida en el mejor de los casos. Es aproximadamente 3 veces más rápido en el peor de los casos.

Si está buscando el primer elemento distinto de cero, puede usar el siguiente truco:

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

Es muy rapido solución "pura y pura" , pero falla en algunos casos que se describen a continuación.

La solución aprovecha el hecho de que prácticamente todas las representaciones de cero para tipos numéricos constan de 0bytes. También se aplica a Numpy bool. En versiones recientes de numpy, la argmax()función usa lógica de cortocircuito al procesar el booltipo. La talla debool es 1 byte.

Entonces uno necesita:

• crear una vista de la matriz como bool . No se crea ninguna copia
• utilizar argmax() para encontrar el primer byte distinto de cero mediante lógica de cortocircuito
• recalcular el desplazamiento de este byte al índice del primer elemento distinto de cero por división entera (operador //) del desplazamiento por un tamaño de un solo elemento expresado en bytes (x.itemsize )
• compruebe si en x[idx]realidad es distinto de cero para identificar el caso en el que no hay ningún distinto de cero

Hice un punto de referencia contra la solución numba y lo construí 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')

El resultado en mi máquina es:

---- 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 solución es un 33% más rápida que numba y es "numpy-pure".

Las desventajas:

• no funciona para muchos tipos aceptables como object
• falla por cero negativo que ocasionalmente aparece en floato doublecálculos

Como usuario de matlab desde hace mucho tiempo, he estado buscando una solución eficiente a este problema durante bastante tiempo. Finalmente, motivado por discusiones y proposiciones en este hilo , he tratado de encontrar una solución que implemente una API similar a la que se sugirió aquí , admitiendo por el momento solo matrices 1D.

Lo usarías así

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)

Los operadores de condición admitidos son: cmp_equal, cmp_not_equal, cmp_larger, cmp_smaller, cmp_larger_eq, cmp_smaller_eq. Por razones de eficiencia, la extensión está escrita en c.

Aquí encontrará la fuente, los puntos de referencia y otros detalles:

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

para el uso en nuestro equipo (anaconda en linux y macos) he hecho un instalador anaconda que simplifica la instalación, puede usarlo como se describe aquí

https://anaconda.org/roebel/py_find_1st

Solo tenga en cuenta que si está haciendo una secuencia de búsquedas, la ganancia de rendimiento al hacer algo inteligente como convertir a una cadena, podría perderse en el ciclo externo si la dimensión de búsqueda no es lo suficientemente grande. Vea cómo el rendimiento de iterar find1 que usa el truco de conversión de cadenas propuesto anteriormente y find2 que usa argmax a lo largo del eje interior (más un ajuste para garantizar que no coincida devuelve -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)

salidas

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

Dicho esto, un hallazgo escrito en C sería al menos un poco más rápido que cualquiera de estos enfoques

Qué tal esto

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

Puede convertir su matriz en a listy usar su index()método:

i = list(array).index(item)

Hasta donde yo sé, este es un método compilado en C.