У меня есть одномерный массив в numpy, и я хочу найти позицию индекса, в которой значение превышает значение в массиве numpy.

Например

aa = range(-10,10)

Найдите позицию, в aaкоторой значение 5будет превышено.

Это немного быстрее (и выглядит лучше)

np.argmax(aa>5)

Поскольку argmaxостановится на первом True(«В случае нескольких повторений максимальных значений, возвращаются индексы, соответствующие первому вхождению.») И не сохраняет другой список.

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

учитывая отсортированное содержимое вашего массива, есть еще более быстрый метод: searchsorted .

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

Меня это тоже интересовало, и я сравнил все предложенные ответы с perfplot . (Отказ от ответственности: я являюсь автором perfplot.)

Если вы знаете, что просматриваемый массив уже отсортирован , тогда

numpy.searchsorted(a, alpha)

это для вас. Это операция O (log (n)), т.е. скорость практически не зависит от размера массива. Вы не можете быть быстрее этого.

Если вы ничего не знаете о своем массиве, вы не ошибетесь с

numpy.argmax(a > alpha)

Уже отсортировано:

Несортированный:

Код для воспроизведения сюжета:

import numpy
import perfplot


alpha = 0.5
numpy.random.seed(0)


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)


perfplot.save(
    "out.png",
    # 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, 23)],
    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

Массивы с постоянным шагом между элементами

В случае a rangeили любого другого линейно возрастающего массива вы можете просто вычислить индекс программно, без необходимости фактически перебирать массив:

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

Возможно, это можно было бы немного улучшить. Я убедился, что он работает правильно для нескольких образцов массивов и значений, но это не значит, что там не может быть ошибок, особенно учитывая, что он использует числа с плавающей запятой ...

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

Учитывая, что он может вычислить позицию без какой-либо итерации, он будет постоянным time ( O(1)) и, вероятно, сможет превзойти все другие упомянутые подходы. Однако для этого требуется постоянный шаг в массиве, иначе это приведет к неправильным результатам.

Общее решение с использованием numba

Более общий подход будет использовать функцию numba:

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

Это будет работать для любого массива, но он должен перебирать массив, поэтому в среднем случае это будет O(n):

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

Контрольный показатель

Несмотря на то, что Нико Шлёмер уже предоставил несколько тестов, я подумал, что было бы полезно включить мои новые решения и протестировать их на разные «значения».

Схема тестирования:

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

и графики были созданы с использованием:

%matplotlib notebook
b.plot()

элемент находится в начале

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

Лучше всего работает функция numba, за которой следуют функция вычисления и функция поиска. Остальные решения работают намного хуже.

элемент находится в конце

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

Для небольших массивов функция numba работает потрясающе быстро, однако для массивов большего размера она уступает функции calculate-function и searchsorted.

элемент находится в 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")

Это более интересно. Снова numba и функция вычисления работают отлично, однако на самом деле это вызывает наихудший случай сортировки по поиску, который в данном случае действительно не работает.

Сравнение функций при отсутствии значения, удовлетворяющего условию

Еще один интересный момент: как ведут себя эти функции, если нет значения, индекс которого должен быть возвращен:

arr = np.ones(100)
value = 2

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

С таким результатом:

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 и numba просто возвращают неверное значение. Однако searchsortedи numbaвернуть индекс, который не является допустимым индексом для массива.

Функции where, min, nonzeroи calculateбросить исключение. Однако только исключение calculateговорит что-то полезное.

Это означает, что на самом деле нужно обернуть эти вызовы в соответствующую функцию-оболочку, которая перехватывает исключения или недопустимые возвращаемые значения и обрабатывает их соответствующим образом, по крайней мере, если вы не уверены, может ли значение находиться в массиве.

Примечание. Параметры расчета и searchsortedработают только в особых условиях. Функция «вычислить» требует постоянного шага, а функция searchsorted требует сортировки массива. Таким образом, они могут быть полезны в определенных обстоятельствах, но не являются общими решениями этой проблемы. Если вы имеете дело с отсортированными списками Python, вы можете взглянуть на модуль bisect вместо использования Numpys searchsorted.

Я хотел бы предложить

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

Это вернет наименьший индекс, в котором выполняется условие, и бесконечность, если условие никогда не выполняется (и whereвозвращает пустой массив).

Я бы пошел с

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

где V- вектор (1d массив), x- значение и i- результирующий индекс.