Массивы с постоянным шагом между элементами
В случае 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
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]
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.
