Eu tenho uma matriz 1D em numpy e quero encontrar a posição do índice em que um valor excede o valor na matriz numpy.

Por exemplo

aa = range(-10,10)

Encontre a posição em aaque o valor 5é excedido.

Isso é um pouco mais rápido (e parece melhor)

np.argmax(aa>5)

Desde argmaxque parará na primeira True("No caso de várias ocorrências dos valores máximos, os índices correspondentes à primeira ocorrência serão retornados.") E não salva outra lista.

In [2]: N = 10000

In [3]: aa = np.arange(-N,N)

In [4]: timeit np.argmax(aa>N/2)
100000 loops, best of 3: 52.3 us per loop

In [5]: timeit np.where(aa>N/2)[0][0]
10000 loops, best of 3: 141 us per loop

In [6]: timeit np.nonzero(aa>N/2)[0][0]
10000 loops, best of 3: 142 us per loop

dado o conteúdo classificado da sua matriz, existe um método ainda mais rápido: a classificação da pesquisa .

import time
N = 10000
aa = np.arange(-N,N)
%timeit np.searchsorted(aa, N/2)+1
%timeit np.argmax(aa>N/2)
%timeit np.where(aa>N/2)[0][0]
%timeit np.nonzero(aa>N/2)[0][0]

# Output
100000 loops, best of 3: 5.97 µs per loop
10000 loops, best of 3: 46.3 µs per loop
10000 loops, best of 3: 154 µs per loop
10000 loops, best of 3: 154 µs per loop

Eu também estava interessado nisso e comparei todas as respostas sugeridas com o perfplot . (Aviso: sou o autor do perfplot.)

Se você sabe que a matriz que você está visualizando já está classificada ,

numpy.searchsorted(a, alpha)

é para você. É uma operação de tempo constante, ou seja, a velocidade não depende do tamanho da matriz. Você não pode ficar mais rápido do que isso.

Se você não sabe nada sobre sua matriz, não está errado com

numpy.argmax(a > alpha)

Já classificado:

Não triados:

Código para reproduzir o gráfico:

import numpy
import perfplot


alpha = 0.5

def argmax(data):
    return numpy.argmax(data > alpha)

def where(data):
    return numpy.where(data > alpha)[0][0]

def nonzero(data):
    return numpy.nonzero(data > alpha)[0][0]

def searchsorted(data):
    return numpy.searchsorted(data, alpha)

out = perfplot.show(
    # setup=numpy.random.rand,
    setup=lambda n: numpy.sort(numpy.random.rand(n)),
    kernels=[
        argmax, where,
        nonzero,
        searchsorted
        ],
    n_range=[2**k for k in range(2, 20)],
    logx=True,
    logy=True,
    xlabel='len(array)'
    )

In [34]: a=np.arange(-10,10)

In [35]: a
Out[35]:
array([-10,  -9,  -8,  -7,  -6,  -5,  -4,  -3,  -2,  -1,   0,   1,   2,
         3,   4,   5,   6,   7,   8,   9])

In [36]: np.where(a>5)
Out[36]: (array([16, 17, 18, 19]),)

In [37]: np.where(a>5)[0][0]
Out[37]: 16

Matrizes que possuem um passo constante entre elementos

No caso de uma rangeou qualquer outra matriz de aumento linear, você pode simplesmente calcular o índice programaticamente, sem necessidade de iterar a matriz:

def first_index_calculate_range_like(val, arr):
    if len(arr) == 0:
        raise ValueError('no value greater than {}'.format(val))
    elif len(arr) == 1:
        if arr[0] > val:
            return 0
        else:
            raise ValueError('no value greater than {}'.format(val))

    first_value = arr[0]
    step = arr[1] - first_value
    # For linearly decreasing arrays or constant arrays we only need to check
    # the first element, because if that does not satisfy the condition
    # no other element will.
    if step <= 0:
        if first_value > val:
            return 0
        else:
            raise ValueError('no value greater than {}'.format(val))

    calculated_position = (val - first_value) / step

    if calculated_position < 0:
        return 0
    elif calculated_position > len(arr) - 1:
        raise ValueError('no value greater than {}'.format(val))

    return int(calculated_position) + 1

Provavelmente, poderia-se melhorar um pouco. Verifiquei se ele funciona corretamente para algumas matrizes e valores de amostra, mas isso não significa que não possam haver erros, especialmente considerando que ele usa flutuadores ...

>>> import numpy as np
>>> first_index_calculate_range_like(5, np.arange(-10, 10))
16
>>> np.arange(-10, 10)[16]  # double check
6

>>> first_index_calculate_range_like(4.8, np.arange(-10, 10))
15

Dado que ele pode calcular a posição sem nenhuma iteração, será um tempo constante ( O(1)) e provavelmente poderá superar todas as outras abordagens mencionadas. No entanto, requer uma etapa constante na matriz, caso contrário, produzirá resultados errados.

Solução geral usando numba

Uma abordagem mais geral seria usar uma função numba:

@nb.njit
def first_index_numba(val, arr):
    for idx in range(len(arr)):
        if arr[idx] > val:
            return idx
    return -1

Isso funcionará para qualquer matriz, mas precisará iterar sobre a matriz; portanto, no caso médio, será O(n):

>>> first_index_numba(4.8, np.arange(-10, 10))
15
>>> first_index_numba(5, np.arange(-10, 10))
16

Referência

Embora Nico Schlömer já tenha fornecido algumas referências, achei que seria útil incluir minhas novas soluções e testar diferentes "valores".

A configuração do teste:

import numpy as np
import math
import numba as nb

def first_index_using_argmax(val, arr):
    return np.argmax(arr > val)

def first_index_using_where(val, arr):
    return np.where(arr > val)[0][0]

def first_index_using_nonzero(val, arr):
    return np.nonzero(arr > val)[0][0]

def first_index_using_searchsorted(val, arr):
    return np.searchsorted(arr, val) + 1

def first_index_using_min(val, arr):
    return np.min(np.where(arr > val))

def first_index_calculate_range_like(val, arr):
    if len(arr) == 0:
        raise ValueError('empty array')
    elif len(arr) == 1:
        if arr[0] > val:
            return 0
        else:
            raise ValueError('no value greater than {}'.format(val))

    first_value = arr[0]
    step = arr[1] - first_value
    if step <= 0:
        if first_value > val:
            return 0
        else:
            raise ValueError('no value greater than {}'.format(val))

    calculated_position = (val - first_value) / step

    if calculated_position < 0:
        return 0
    elif calculated_position > len(arr) - 1:
        raise ValueError('no value greater than {}'.format(val))

    return int(calculated_position) + 1

@nb.njit
def first_index_numba(val, arr):
    for idx in range(len(arr)):
        if arr[idx] > val:
            return idx
    return -1

funcs = [
    first_index_using_argmax, 
    first_index_using_min, 
    first_index_using_nonzero,
    first_index_calculate_range_like, 
    first_index_numba, 
    first_index_using_searchsorted, 
    first_index_using_where
]

from simple_benchmark import benchmark, MultiArgument

e as parcelas foram geradas usando:

%matplotlib notebook
b.plot()

item está no começo

b = benchmark(
    funcs,
    {2**i: MultiArgument([0, np.arange(2**i)]) for i in range(2, 20)},
    argument_name="array size")

A função numba apresenta o melhor desempenho, seguida pela função de cálculo e pela função de busca variada. As outras soluções apresentam desempenho muito pior.

item está no final

b = benchmark(
    funcs,
    {2**i: MultiArgument([2**i-2, np.arange(2**i)]) for i in range(2, 20)},
    argument_name="array size")

Para matrizes pequenas, a função numba executa incrivelmente rápido, no entanto, para matrizes maiores, ela é superada pela função de cálculo e pela função de seleção de pesquisa.

item está no sqrt (len)

b = benchmark(
    funcs,
    {2**i: MultiArgument([np.sqrt(2**i), np.arange(2**i)]) for i in range(2, 20)},
    argument_name="array size")

Isso é mais interessante. Novamente, o numba e a função de cálculo têm um ótimo desempenho, no entanto, isso está realmente desencadeando o pior caso de searchsorted, o que realmente não funciona bem nesse caso.

Comparação das funções quando nenhum valor satisfaz a condição

Outro ponto interessante é como essas funções se comportam se não houver valor cujo índice deva ser retornado:

arr = np.ones(100)
value = 2

for func in funcs:
    print(func.__name__)
    try:
        print('-->', func(value, arr))
    except Exception as e:
        print('-->', e)

Com este resultado:

first_index_using_argmax
--> 0
first_index_using_min
--> zero-size array to reduction operation minimum which has no identity
first_index_using_nonzero
--> index 0 is out of bounds for axis 0 with size 0
first_index_calculate_range_like
--> no value greater than 2
first_index_numba
--> -1
first_index_using_searchsorted
--> 101
first_index_using_where
--> index 0 is out of bounds for axis 0 with size 0

Searchsorted, argmax e numba simplesmente retornam um valor errado. No entanto searchsortede numbaretornar um índice que não é um índice válido para a matriz.

As funções where, min, nonzeroe calculatelançar uma exceção. No entanto, apenas a exceção para calculaterealmente diz algo útil.

Isso significa que é necessário agrupar essas chamadas em uma função de wrapper apropriada que captura exceções ou valores de retorno inválidos e manipula-os adequadamente, pelo menos se você não tiver certeza se o valor pode estar na matriz.

Nota: O cálculo e as searchsortedopções funcionam apenas em condições especiais. A função "calcular" requer uma etapa constante e a pesquisa ordenada exige que a matriz seja classificada. Portanto, eles podem ser úteis nas circunstâncias certas, mas não são soluções gerais para esse problema. Caso esteja lidando com listas Python classificadas, você pode dar uma olhada no módulo bisect em vez de usar o Numpys searchsorted.

Eu gostaria de propor

np.min(np.append(np.where(aa>5)[0],np.inf))

Isso retornará o menor índice em que a condição for atendida, enquanto retornará o infinito se a condição nunca for atendida (e whereretorna uma matriz vazia).

Eu iria com

i = np.min(np.where(V >= x))

onde Vé vetor (matriz 1d), xé o valor e ié o índice resultante.