Numpy: encontre o primeiro índice de valor rapidamente

Como posso encontrar o índice da primeira ocorrência de um número em uma matriz Numpy? A velocidade é importante para mim. Não estou interessado nas seguintes respostas porque elas examinam todo o array e não param quando encontram a primeira ocorrência:

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

Nota 1: nenhuma das respostas dessa pergunta parece relevante. Existe uma função Numpy para retornar o primeiro índice de algo em um array?

Nota 2: usar um método compilado em C é preferível a um loop Python.

Há uma solicitação de recurso agendada para Numpy 2.0.0: https://github.com/numpy/numpy/issues/2269

Embora seja tarde demais para você, mas para referência futura: Usar numba ( 1 ) é a maneira mais fácil até que numpy o implemente. Se você usa a distribuição anaconda python, ela já deve estar instalada. O código será compilado para que seja 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

e depois:

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

Eu fiz uma referência para vários métodos:

• argwhere
• nonzero como na pergunta
• .tostring() como na resposta de @Rob Reilink
• loop de python
• Loop Fortran

Os códigos Python e Fortran estão disponíveis. Eu pulei os pouco promissores, como converter para uma lista.

Os resultados em escala logarítmica. O eixo X é a posição da agulha (leva mais tempo para descobrir se está mais abaixo na matriz); o último valor é uma agulha que não está na matriz. O eixo Y é a hora de encontrá-lo.

O array tinha 1 milhão de elementos e os testes foram executados 100 vezes. Os resultados ainda variam um pouco, mas a tendência qualitativa é clara: o Python e o f2py fecham no primeiro elemento, então eles escalam de forma diferente. Python fica muito lento se a agulha não estiver nos primeiros 1%, enquanto f2pyé rápido (mas você precisa compilá-lo).

Para resumir, f2py é a solução mais rápida , especialmente se o ponteiro aparecer bem cedo.

Não é embutido, o que é irritante, mas na verdade é apenas 2 minutos de trabalho. Adicione isso a um arquivo chamado 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 você estiver procurando por algo diferente integer, basta alterar o tipo. Em seguida, compile usando:

f2py -c -m search search.f90

depois disso, você pode fazer (do Python):

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

Você pode converter uma matriz booleana em uma string Python usando array.tostring()e, em seguida, usando o método find ():

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

No entanto, isso envolve a cópia dos dados, já que as strings do Python precisam ser imutáveis. Uma vantagem é que você também pode pesquisar, por exemplo, uma borda ascendente, encontrando\x00\x01

Em caso de np.searchsortedtrabalhos de matrizes classificadas .

Acho que você encontrou um problema em que um método diferente e algum conhecimento a priori do array realmente ajudariam. O tipo de coisa em que você tem uma probabilidade X de encontrar sua resposta nos primeiros Y por cento dos dados. A divisão do problema com a esperança de ter sorte e então fazer isso em python com uma compreensão de lista aninhada ou algo assim.

Escrever uma função C para fazer essa força bruta também não é muito difícil usando ctypes .

O código C que hackeado junto (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 o 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))

e recebo 92.

Enrole o python em uma função adequada e pronto.

A versão C é muito (~ 20x) mais rápida para este seed (avisando que não sou bom com o 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 já apresentou uma numbafunção para encontrar o primeiro índice, mas que só funciona para arrays 1D. Com np.ndenumeratevocê também pode encontrar o primeiro índice em uma matriz de dimensão arbitrária:

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 amostra:

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

Os tempos mostram que é semelhante em desempenho à solução 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 sua lista estiver ordenada , você pode realizar uma busca muito rápida de índice com o pacote 'bisect'. É O (log (n)) em vez de O (n).

bisect.bisect(a, x)

encontra x no array a, definitivamente mais rápido no caso classificado do que qualquer rotina C passando por todos os primeiros elementos (para listas longas o suficiente).

É bom saber às vezes.

Até onde eu sei, apenas np.any e np.all em matrizes booleanas estão em curto-circuito.

No seu caso, numpy tem que percorrer todo o array duas vezes, uma para criar a condição booleana e uma segunda vez para encontrar os índices.

Minha recomendação neste caso seria usar cíton. Acho que deve ser fácil ajustar um exemplo para este caso, especialmente se você não precisa de muita flexibilidade para diferentes tipos e formatos.

Eu precisava disso para o meu trabalho, então aprendi a interface C do Python e do Numpy e escrevi minha própria. http://pastebin.com/GtcXuLyd É apenas para arrays 1-D, mas funciona para a maioria dos tipos de dados (int, float ou strings) e os testes mostraram que é novamente cerca de 20 vezes mais rápido do que a abordagem esperada em Python- puro entorpecido.

Este problema pode ser resolvido de forma eficaz em puro numpy, processando a matriz em blocos:

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

A matriz é processada em pedaços de tamanho step. Quanto stepmais longa a etapa, mais rápido é o processamento da matriz zerada (pior caso). Quanto menor for, mais rápido será o processamento da matriz com um valor diferente de zero no início. O truque é começar com um valor pequeno stepe aumentá-lo exponencialmente. Além disso, não há necessidade de incrementá-lo acima de algum limite devido aos benefícios limitados.

Eu comparei a solução com a solução ndarary.nonzero e numba pura contra 10 milhões de conjuntos de flutuadores.

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 os resultados na minha 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.nonzeroé definitivamente mais solto. A solução numba é cerca de 5 vezes mais rápida para o melhor caso. É cerca de 3 vezes mais rápido no pior dos casos.

Se você está procurando o primeiro elemento diferente de zero, pode usar o seguinte hack:

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

É uma solução "numpy-pura" muito rápida , mas falha em alguns casos discutidos abaixo.

A solução tira vantagem do fato de que praticamente toda representação de zero para tipos numéricos consiste em 0bytes. Isso se aplica a numpy's booltambém. Em versões recentes de numpy, a argmax()função usa lógica de curto-circuito ao processar o booltipo. O tamanho de boolé 1 byte.

Então, é preciso:

• crie uma visualização da matriz como bool. Nenhuma cópia é criada
• use argmax()para encontrar o primeiro byte diferente de zero usando lógica de curto-circuito
• recalcular o deslocamento deste byte para o índice do primeiro elemento diferente de zero por divisão inteira (operador //) do deslocamento por um tamanho de um único elemento expresso em bytes ( x.itemsize)
• verifique se x[idx]é realmente diferente de zero para identificar o caso quando nenhum diferente de zero está presente

Eu fiz alguns benchmarks contra a solução numba e a 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')

Os resultados em minha máquina são:

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

A solução é 33% mais rápida do que numba e é "numpy-pure".

As desvantagens:

• não funciona para tipos numpy aceitáveis ​​como object
• falha para zero negativo que ocasionalmente aparece em floatou doublecálculos

Como usuário de matlab de longa data, há muito tempo procuro uma solução eficiente para esse problema. Finalmente, motivado por discussões sobre as proposições neste tópico , tentei encontrar uma solução que implementasse uma API semelhante à sugerida aqui , suportando por enquanto apenas matrizes 1D.

Você usaria assim

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)

Os operadores de condição suportados são: cmp_equal, cmp_not_equal, cmp_larger, cmp_smaller, cmp_larger_eq, cmp_smaller_eq. Para eficiência, a extensão é escrita em c.

Você encontra a fonte, benchmarks e outros detalhes aqui:

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

para uso em nossa equipe (anaconda no linux e macos) eu fiz um instalador anaconda que simplifica a instalação, você pode usá-lo conforme descrito aqui

https://anaconda.org/roebel/py_find_1st

Apenas uma observação que se você estiver fazendo uma sequência de pesquisas, o ganho de desempenho de fazer algo inteligente como converter para string pode ser perdido no loop externo se a dimensão da pesquisa não for grande o suficiente. Veja como o desempenho de iteração de find1 que usa o truque de conversão de string proposto acima e find2 que usa argmax ao longo do eixo interno (mais um ajuste para garantir que uma não correspondência retorne como -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)

saídas

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

Dito isso, um achado escrito em C seria pelo menos um pouco mais rápido do que qualquer uma dessas abordagens

que tal agora

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

Você pode converter sua matriz em um liste usar seu index()método:

i = list(array).index(item)

Pelo que eu sei, este é um método compilado em C.