Эффективный поиск Numpy в немонотонном массиве

Question

Я пытаюсь провести что-то похожее на searchsorted, но в случае, когда массив не является полностью монотонным.Скажем, у меня есть скаляр c и одномерный массив x, я хочу найти индексы i всех элементов, такие что x[i] < c <= x[i + 1].Важно отметить, что x не является полностью монотонным.

Следующий код работает, но я просто хотел бы знать, является ли это наиболее эффективным способом сделать это, или есть более простой способ:

x = np.array([1,2,3,1,2,3,1,2,3])
c = 2.5
t = c > x[:-1]
u = c <= x[1:]
v = t*u
i = v.nonzero()[0]

Или в одной строке кода:

i = ( (c > x[:-1]) * (c <= x[1:] ).nonzero()[0]

Является ли это наиболее эффективным способом восстановления этих индексов?

Два дополнительных вопроса.

1. Существует ли простой способ расширить это до случая, когда c - это одномерный массив, а x - это двумерный массив, где c имеет столько элементов, сколько "строк "в x, и я выполняю этот поиск для каждого элемента c в соответствующей" строке "из x?

2. Моя конечная цель - сделать это стрехмерный случай.То есть, предположим, что c все еще является одномерным вектором с n элементами.Теперь пусть x будет трехмерным массивом с размерами j на n на k.Есть ли способ сделать №1 выше для каждой «подматрицы» в x?В основном, выполняя # 1 выше j раз.

Например:

x1 = np.array([1,2,3,1,2,3],[1,2,3,1,2,3],[1,2,3,1,2,3])
x2 = x1 + 1
x = np.array([x1,x2])
c = np.array([1.5,2.5,3.5])

Под # 1 выше, когда мы сравниваем c и x1, мы получим: [[0,4],[1,5],[]]

Когда мы сравним c и x2, мы получим: [[],[0,4],[1,5]]

Наконец, под # 2 я бы хотел получить:

[[[0,4],[1,5],[]],
 [[],[0,4],[1,5]]]

Answer 1

В такой задаче numpy делает слишком много сравнений.Вы можете выиграть в 5 раз с Numba.Нет трудностей в адаптации для 3-х измерений.

@numba.njit
def ind(x,c):
    res = empty_like(x)
    i=j=0
    while i < x.size-1:
        if x[i]<c and c<=x[i+1]:
            res[j]=i
            j+=1
        i+=1 
    return res[:j]

Answer 2

Мы могли бы сравнить один раз, чтобы дать нам логическую маску, и повторно использовать ее с отрицанием, чтобы получить другой массив сравнения, а также использовать нарезку -

m = c > x
i = np.flatnonzero( m[:-1] & ~m[1:] )

Мы можем расширить ее до x как 2D и c как 1D случай с циклом, но сделайте с ним минимальные вычисления, предварительно вычислив генерацию масок в векторизованном виде, например:

m = c[:,None] > x
m2 = m[:,:-1] & ~m[:,1:]
i = [np.flatnonzero( mi ) for mi in m2]