Numpy: быстро найти первый индекс значения

Как я могу найти индекс первого вхождения числа в массиве Numpy? Для меня важна скорость. Меня не интересуют следующие ответы, потому что они сканируют весь массив и не останавливаются, когда находят первое вхождение:

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

Примечание 1. Ни один из ответов на этот вопрос не кажется релевантным. Существует ли функция Numpy для возврата первого индекса чего-либо в массиве?

Примечание 2: использование C-скомпилированного метода предпочтительнее цикла Python.

Для Numpy 2.0.0 есть запрос функции: https://github.com/numpy/numpy/issues/2269

Хотя для вас уже слишком поздно, но для справки в будущем: использование numba ( 1 ) - самый простой способ, пока numpy не реализует его. Если вы используете дистрибутив anaconda python, он уже должен быть установлен. Код будет скомпилирован, так что все будет быстро.

@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

а потом:

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

Я проверил несколько методов:

• argwhere
• nonzero как в вопросе
• .tostring() как в ответе @Rob Reilink
• цикл Python
• Цикл Fortran

Python и Fortran кода доступны. Я пропустил бесперспективные, например преобразование в список.

Результаты в логарифмическом масштабе. Ось X - это положение стрелки (требуется больше времени, чтобы определить, находится ли она дальше по массиву); последнее значение - игла, которой нет в массиве. Ось Y - время найти его.

В массиве 1 миллион элементов, и тесты выполнялись 100 раз. Результаты все еще немного колеблются, но качественная тенденция очевидна: Python и f2py завершают работу на первом элементе, поэтому масштабируются по-разному. Python становится слишком медленным, если стрелка не находится в первых 1%, тогда как он f2pyработает быстро (но вам нужно его скомпилировать).

Подводя итог, f2py - самое быстрое решение , особенно если игла появляется довольно рано.

Он не встроен, что раздражает, но на самом деле это всего 2 минуты работы. Добавьте это в файл с именем 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

Если вы ищете что-то другое integer, просто измените тип. Затем скомпилируйте, используя:

f2py -c -m search search.f90

после чего вы можете сделать (из Python):

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

Вы можете преобразовать логический массив в строку Python, используя, array.tostring()а затем используя метод find ():

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

Однако это подразумевает копирование данных, поскольку строки Python должны быть неизменными. Преимущество состоит в том, что вы также можете искать, например, нарастающий фронт, найдя\x00\x01

В случае сортированных массивов np.searchsortedработает.

Я думаю, вы столкнулись с проблемой, когда действительно помогли бы другой метод и некоторое априорное знание массива. То, что у вас есть X вероятность найти свой ответ в первых Y процентах данных. Разделение проблемы с надеждой на удачу, а затем выполнение этого на python с пониманием вложенного списка или что-то в этом роде.

Написание функции C для выполнения этой грубой силы также не так уж сложно с использованием ctypes .

Код C, который я взломал вместе (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);
}

и питон:

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

и я получаю 92.

Оберните питон в подходящую функцию, и готово.

Версия C намного (~ 20x) быстрее для этого семени (предупреждение, я не очень хорошо разбираюсь в 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 уже представил numbaфункцию для поиска первого индекса, но она работает только для одномерных массивов. С помощью np.ndenumerateвы также можете найти первый индекс в массиве произвольной размерности:

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

Пример кейса:

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

Тайминги показывают, что по производительности он похож на решение 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

Если ваш список отсортирован , вы можете добиться очень быстрого поиска по индексу с помощью пакета 'bisect'. Это O (log (n)) вместо O (n).

bisect.bisect(a, x)

находит x в массиве a, что определенно быстрее в отсортированном случае, чем любая C-подпрограмма, проходящая через все первые элементы (для достаточно длинных списков).

Иногда полезно знать.

Насколько мне известно, закорочены только np.any и np.all на булевых массивах.

В вашем случае numpy должен дважды пройти через весь массив: один раз для создания логического условия и второй раз для поиска индексов.

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

Мне это было нужно для работы, поэтому я изучил Python и интерфейс C Numpy и написал свой собственный. http://pastebin.com/GtcXuLyd Это только для одномерных массивов, но работает для большинства типов данных (int, float или strings), и тестирование показало, что он снова примерно в 20 раз быстрее, чем ожидаемый подход в чистом Python- тупой.

Эта проблема может быть эффективно решена в чистом numpy путем обработки массива кусками:

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

Массив обрабатывается по размеру step. Чем stepдлиннее шаг, тем быстрее выполняется обработка обнуленного массива (худший случай). Чем он меньше, тем быстрее обрабатывается массив с ненулевым значением в начале. Уловка состоит в том, чтобы начать с малого stepи увеличивать его экспоненциально. Более того, нет необходимости увеличивать его выше некоторого порога из-за ограниченных преимуществ.

Я сравнил решение с чистым решением ndarary.nonzero и numba с 10 миллионами массивов с плавающей запятой.

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

И результаты на моей машине:

---- 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определенно слабее. Решение numba в лучшем случае примерно в 5 раз быстрее. В худшем случае это примерно в 3 раза быстрее.

Если вы ищете первый ненулевой элемент, вы можете использовать следующий прием:

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

Это очень быстрое "безупречное" решение, но оно не работает в некоторых случаях, обсуждаемых ниже.

Решение использует тот факт, что почти все представление нуля для числовых типов состоит из 0байтов. Это относится и к numpy bool. В последних версиях numpy argmax()функция использует логику короткого замыкания при обработке boolтипа. Размер bool1 байт.

Итак, нужно:

• создать представление массива как bool. Копия не создается
• использовать argmax()для поиска первого ненулевого байта с помощью логики короткого замыкания
• пересчитать смещение этого байта в индекс первого ненулевого элемента путем целочисленного деления (оператор //) смещения на размер одного элемента, выраженного в байтах ( x.itemsize)
• проверьте, x[idx]действительно ли ненулевое значение, чтобы определить случай, когда ненулевое значение отсутствует

Я сделал несколько тестов против решения numba и построил его 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')

Результат на моей машине:

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

Решение на 33% быстрее, чем numba, и оно "безупречно".

Недостатки:

• не работает для приемлемых типов, таких как object
• не выполняется для отрицательного нуля , что время от времени появляется в floatили doubleвычислении

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

Вы бы использовали это так

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)

Поддерживаются следующие операторы условий: cmp_equal, cmp_not_equal, cmp_larger, cmp_smaller, cmp_larger_eq, cmp_smaller_eq. Для эффективности расширение написано в c.

Вы можете найти исходный код, тесты и другие подробности здесь:

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

для использования в нашей команде (анаконда на linux и macos) Я сделал установщик анаконды, который упрощает установку, вы можете использовать его, как описано здесь

https://anaconda.org/roebel/py_find_1st

Замечу, что если вы выполняете последовательность поисков, выигрыш в производительности от таких умных действий, как преобразование в строку, может быть потерян во внешнем цикле, если размер поиска недостаточно велик. Посмотрите, как производительность итерации find1, использующей предложенный выше трюк с преобразованием строк, и find2, использующей argmax вдоль внутренней оси (плюс корректировка, гарантирующая, что несоответствие возвращается как -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)

выходы

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

Тем не менее, находка, написанная на C, будет по крайней мере немного быстрее, чем любой из этих подходов.

как насчет этого

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

Вы можете скрыть свой массив в listи использовать его index()метод:

i = list(array).index(item)

Насколько мне известно, это метод, скомпилированный на C.