Эффективно вернуть индекс первого значения, удовлетворяющего условию в массиве

numba

С numba можно оптимизировать сценарии оба. Синтаксически вам нужно только построить функцию с помощью простого цикла for:

from numba import njit

@njit
def get_first_index_nb(A, k):
    for i in range(len(A)):
        if A[i] > k:
            return i
    return -1

idx = get_first_index_nb(A, 0.9)

Numba повышает производительность за счет JIT ("Just In Time") компиляции кода и использования Оптимизация на уровне ЦП. Цикл обычныйfor без декоратора @njit обычно будет помедленнее, чем методы, которые вы уже пробовали, для случая, когда условие выполняется поздно.

Для числовой серии df['data'] Pandas вы можете просто передать представление NumPy в JIT-скомпилированную функцию:

idx = get_first_index_nb(df['data'].values, 0.9)

Обобщение

Поскольку numba разрешает функции как аргументы, и предполагая, что переданная функция также может быть JIT-скомпилирована, вы можете получить метод для вычисления индекса пth, где выполняется условие для произвольного func.

@njit
def get_nth_index_count(A, func, count):
    c = 0
    for i in range(len(A)):
        if func(A[i]):
            c += 1
            if c == count:
                return i
    return -1

@njit
def func(val):
    return val > 0.9

# get index of 3rd value where func evaluates to True
idx = get_nth_index_count(arr, func, 3)

Для 3-го значения последний вы можете передать обратное, arr[::-1], и отрицать результат len(arr) - 1, - 1, необходимый для учета 0-индексации.

Тестирование производительности

# Python 3.6.5, NumPy 1.14.3, Numba 0.38.0

np.random.seed(0)
arr = np.random.rand(10**7)
m = 0.9
n = 0.999999

@njit
def get_first_index_nb(A, k):
    for i in range(len(A)):
        if A[i] > k:
            return i
    return -1

def get_first_index_np(A, k):
    for i in range(len(A)):
        if A[i] > k:
            return i
    return -1

%timeit get_first_index_nb(arr, m)                                 # 375 ns
%timeit get_first_index_np(arr, m)                                 # 2.71 µs
%timeit next(iter(np.where(arr > m)[0]), -1)                       # 43.5 ms
%timeit next((idx for idx, val in enumerate(arr) if val > m), -1)  # 2.5 µs

%timeit get_first_index_nb(arr, n)                                 # 204 µs
%timeit get_first_index_np(arr, n)                                 # 44.8 ms
%timeit next(iter(np.where(arr > n)[0]), -1)                       # 21.4 ms
%timeit next((idx for idx, val in enumerate(arr) if val > n), -1)  # 39.2 ms

Я тоже хотел сделать что-то подобное и обнаружил, что решения, представленные в этом вопросе, мне не очень помогают. В частности, решение numba было для меня намного медленнее, чем более традиционные методы, представленные в самом вопросе. У меня есть список times_all, обычно состоящий из десятков тысяч элементов, и я хочу найти индекс первого элемента times_all, который больше, чем time_event. И у меня тысячи time_event. Мое решение состоит в том, чтобы разделить times_all на куски, например, по 100 элементов, сначала решить, какому временному сегменту принадлежит time_event, сохранить индекс первого элемента этого сегмента, затем найти, какой индекс в этом сегменте, и сложить два индекса. Вот минимальный код. Для меня он работает на порядки быстрее, чем другие решения на этой странице.

def event_time_2_index(time_event, times_all, STEPS=100):
    import numpy as np
    time_indices_jumps = np.arange(0, len(times_all), STEPS)
    time_list_jumps = [times_all[idx] for idx in time_indices_jumps]

    time_list_jumps_idx = next((idx for idx, val in enumerate(time_list_jumps)\
                          if val > time_event), -1)
    index_in_jumps = time_indices_jumps[time_list_jumps_idx-1]
    times_cropped = times_all[index_in_jumps:]
    event_index_rel = next((idx for idx, val in enumerate(times_cropped) \
                      if val > time_event), -1)

    event_index = event_index_rel + index_in_jumps
    return event_index