Numpy: швидко знайдіть перший індекс вартості

Як я можу знайти індекс першого появи числа в масиві Numpy? Швидкість для мене важлива. Наступні відповіді мене не цікавлять, оскільки вони сканують весь масив і не зупиняються, коли знаходять перше явище:

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

Примітка 1: жоден з відповідей із цього питання не видається релевантним Чи існує функція Numpy для повернення першого індексу чогось у масиві?

Примітка 2: використання методу, складеного С, є кращим для циклу Python.

Існує запит на функцію для цього, запланований для Numpy 2.0.0: https://github.com/numpy/numpy/isissue/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 . Я пропустив неперспективних, як перехід до списку.

Результати за журнальною шкалою. Вісь 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 .

Код С, який я зламав разом (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функцію пошуку першого індексу, але він працює лише для 1D масивів. З np.ndenumerateвами також може знайти перший індекс в якості arbitarly одновимірного масиву:

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 повинен пройти весь масив двічі, один раз для створення булевого стану та вдруге для пошуку індексів.

Моєю рекомендацією в цьому випадку буде використання цитону. Я думаю, що слід легко налагодити приклад для цього випадку, особливо якщо вам не потрібна велика гнучкість для різних типів і форм.

Мені це було потрібно для моєї роботи, тому я навчив себе інтерфейс Python та Numpy і написав власний. http://pastebin.com/GtcXuLyd Це лише для 1-D масивів, але працює для більшості типів даних (int, float або string), і тестування показало, що це знову приблизно в 20 разів швидше, ніж очікуваний підхід у чистому Python- онімілий.

Цю проблему можна ефективно вирішити чистим числом, обробивши масив фрагментами:

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

Чистий ndarray.nonzero- це певний невдалий. Рішення numba приблизно в 5 разів швидше для кращого випадку. Це в гіршому випадку приблизно в 3 рази швидше.

Якщо ви шукаєте перший ненульовий елемент, ви можете скористатися наступним хаком:

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

Це дуже швидке "чисте" рішення, але воно не вдається для деяких випадків, обговорених нижче.

Рішення має перевагу з того, що майже все представлення нуля для числових типів складається з 0байтів. Це стосується і numpy's bool. В останніх версіях numpy argmax()функція використовує логіку короткого замикання при обробці boolтипу. Розмір bool- 1 байт.

Отже, потрібно:

• створити подання масиву як 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, подібний до запропонованого тут , підтримуючи на даний момент лише 1D масиви.

Ви б використовували це так

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

для використання в нашій команді (anaconda на Linux та macos) Я зробив інсталятор anaconda, який спрощує встановлення, ви можете використовувати його, як описано тут

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)

Однак, пошук, написаний на мові С, буде хоч трохи швидшим, ніж будь-який із цих підходів

як щодо цього

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

Ви можете перетворити масив у a listта використовувати його index()метод:

i = list(array).index(item)

Наскільки мені відомо, це метод, складений С.