Как мне быстро найти последнее число элементов, соответствующих некоторому критерию, в numpy?

Question

Предположим, у меня есть какой-то массив:

x = numpy.array([1, 2, 3, 4, 5, 1, 1, 1])

И я хочу найти количество последовательных 1 с в конце массива. Один из способов сделать это с помощью цикла:

i = len(x) - 1
while x[i] == 1:
    i = i - 1

Теперь я могу посмотреть на i и определить количество 1 s в конце x. Однако в моем реальном примере x может быть очень большим, как и число 1 с, поэтому я хочу решение, которое:

1. Не использует циклы, а
2. Не пересекает весь массив

Answer 1

Посмотрите на поиск: https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.searchsorted.html#numpy.searchsorted

Используйте его в сочетании с np.where в последнем элементе, чтобы вернуть 0, если! = 1.

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

import numpy as np

x = np.array([1, 2, 3, 4, 5, 1, 1, 1])
y = np.array([1, 2, 3, 4, 5, 1, 1, 1, 4, 5])

# Create a lambda function that accepts x(array)
f = lambda x: np.where(x[-1] != 1, 0, np.searchsorted(x[::-1], 1, side='right'))

print(f(x)) # 3
print(f(y)) # 0

Answer 2

Я согласен с использованием Cython или numba для ускорения цикла и обхода только хвоста массива, но если вы хотите попробовать его с чистым numpy, я бы сказал, что может сделать что-то вроде следующего:

np.argwhere(x[::-1] != 1).ravel()[0]

Переверните массив и возьмите первое не-1 вхождение. Хотя он охватывает весь массив ... так что, вероятно, он не соответствует вашим потребностям.

РЕДАКТИРОВАТЬ: Вот numba подход для полноты

from numba import jit
@jit
def count_trailing_ones(array):
    count = 0
    a = array[::-1]
    for i in range(array.shape[0]):
        if a[i] == 1:
            count += 1
        else:
            return count

Вот эталонный тест, включающий также решения @J ... S и @Kasramvd для массива 800 МБ с парой миллионов трейлинговых. numba очевидно побеждает, но если вы идете на numpy Я бы сказал, @J ... S's argmax является лучшим.

In [102]: %timeit np.argwhere(x[::-1] != 1).ravel()[0]
631 ms ± 1.83 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [103]: %timeit np.argmax(x[::-1] != 1)
117 ms ± 417 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [104]: %timeit kas(x)
915 ms ± 3.41 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [105]: %timeit count_trailing_ones(x)
4.62 ms ± 16.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Answer 3

Если последний элемент всегда равен 1, вы можете сначала обратить массив, а затем использовать argmax, как в

np.argmax(x[::-1]!=1)

Который для данного массива даст

3

Вы можете сначала использовать проверку, чтобы убедиться, что последний элемент имеет значение 1, например

if(x[-1] == 1):
    print(np.argmax(x[::-1]!=1))
else:
    print(0)

Answer 4

Вот один из подходов векторизации:

In [50]: mask = x == 1

In [51]: T_inds = np.where(mask)[0]

In [52]: F_inds = np.where(~mask)[0]

In [53]: last_f_ind = np.where(T_inds[-1] > F_inds)[0][-1]

# x = np.array([1, 2, 3, 4, 5, 1, 1, 1, 4, 5])
In [54]: T_inds[-1] - F_inds[last_f_ind]
Out[54]: 3

Хитрость заключается в том, чтобы найти индекс последнего элемента в не одном элементе, который меньше индекса последнего.

Также обратите внимание, что этот подход будет работать для всех случаев, когда четные не находятся в конце вашего массива (после последних 1 с нет других чисел). Но если вы хотите проверить тот конкретный случай, когда 1 находятся в конце вашего массива, вот более краткий подход:

x.size - np.where(x != 1)[0][-1] - 1
Out[27]: 3

# x != 1 will give you a mask of the indices where their value is not
# equal to one. Then you can use np.where() to find the index of last 
# occurrence of not one items. By subtracting it from size of array you 
# can get the number of consecutive ones.