Нумерованный эквивалент list.index

В низкоуровневой функции, которая называется много раз, мне нужно сделать эквивалент python list.index, но с массивом numpy. Функция должна возвратиться, когда она найдет первое значение, и в противном случае повысит значение ValueError. Что-то вроде:

>>> a = np.array([1, 2, 3])
>>> np_index(a, 1)
0
>>> np_index(a, 10)
Traceback (most recent call last):    
  File "<stdin>", line 1, in <module>
ValueError: 10 not in array

Я хочу, если возможно, избегать цикла Python. np.where не является вариантом, поскольку он всегда выполняет итерацию по всему массиву; Мне нужно что-то, что останавливается после обнаружения первого индекса.

EDIT. Еще одна конкретная информация, связанная с этой проблемой.

  • Примерно в 90% случаев индекс, который я ищу, находится в первых 1/4 до 1/2 массива. Так что здесь потенциально может произойти 2-4 ускорения. Остальные 10% времени не имеют значения в массиве.

  • Я уже профилировал вещи, а вызов np.where - это узкое место, занимающее не менее 50% общей продолжительности выполнения.

  • Не важно, чтобы он поднимал ValueError; он просто должен вернуть что-то, что, очевидно, указывает на то, что значение не находится в массиве.

Я, вероятно, буду кодировать решение в Cython, как было предложено.

Ответы

Ответ 1

См. мой комментарий к вопросу OP для предостережений, но в целом я бы сделал следующее:

import numpy as np
a = np.array([1, 2, 3])
np.min(np.nonzero(a == 2)[0])

Если значение, которое вы ищете, не находится в массиве, вы получите ValueError из-за:

ValueError: zero-size array to ufunc.reduce without identity

потому что вы пытаетесь принять значение min для пустого массива.

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

Ответ 2

Если ваш массив numpy равен 1d массиву, возможно, попробуйте вот так:

a = np.array([1, 2, 3])
print a.tolist().index(2)
>>> 1

Если это не 1d, вы можете найти такой массив, как:

a = np.array([[1, 2, 3],[2,5,6],[0,0,2]])
print a[0,:].tolist().index(2)
>>> 1

print a[1,:].tolist().index(2)
>>> 0

print a[2,:].tolist().index(2)
>>> 2

Ответ 3

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

http://www.scipy.org/Numpy_Example_List_With_Doc#nonzero

В частности, эта часть:

a.nonzero()

Возвращает индексы элементов, отличных от нуля.

Для полной документации обратитесь к numpy.nonzero.

См. также

numpy.nonzero: эквивалентная функция

>>> from numpy import *
>>> y = array([1,3,5,7])
>>> indices = (y >= 5).nonzero()
>>> y[indices]
array([5, 7])
>>> nonzero(y)                                # function also exists
(array([0, 1, 2, 3]),)

Где (http://www.scipy.org/Numpy_Example_List_With_Doc#where) также может вас заинтересовать.

Ответ 4

Вы можете закодировать его в Cython и просто импортировать из Python script. Не нужно переносить весь проект в Cython.

# paste into: indexing.pyx
def index(long[:] lst, long value):
    cdef int i
    for i in range(len(lst)):
        if lst[i] == value:
            return i
    raise ValueError

# import in your .py code
import pyximport
pyximport.install()
from indexing import index

# example
from numpy import zeros
a = zeros(10**6, int)
a[-1] = 1

index(a, 1)
Wall time: 6.07 ms
999999

index(a, 0)
Wall time: 38.1 µs
0

Ответ 5

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

a = numpy.arange(3)
print(list(a).index(2))

>>> 2

Ответ 6

NumPy searchsorted очень похож на индекс списков, за исключением того, что он требует отсортированного массива и ведет себя более численно. Большие различия заключаются в том, что вам не нужно иметь точное соответствие, и вы можете искать, начиная с левой или с правой стороны. Для получения информации о том, как это работает, см. Следующие примеры:

import numpy as np
a = np.array([10, 20, 30])

a.searchsorted(-99) == a.searchsorted(0) == a.searchsorted(10)
# returns index 0 for value 10

a.searchsorted(20.1) == a.searchsorted(29.9) == a.searchsorted(30)
# returns index 2 for value 30

a.searchsorted(30.1) == a.searchsorted(99) == a.searchsorted(np.nan)
# returns index 3 for undefined value

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