Numpy условная проверка с остановкой

Question

Мне нужна твоя помощь. Я хочу пройтись по трехмерному массиву и проверить в одном направлении расстояние между двумя элементами, если оно меньше, значение должно быть True. Как только расстояние становится выше определенного значения, остальные значения в этом измерении должны быть установлены в False.

Вот пример в 1D:

    a = np.array([1,2,2,1,2,5,2,7,1,2])
    b = magic_check_fct(a, threshold=3, axis=0)
    print(b)

   # The expected output is :
   > b = [True, True, True, True, True, False, False, False, False, False]

Для простой проверки результат с a <= threshold будет и не будет ожидаемым результатом:

   > b = [True, True, True, True, True, False, True, False, True, True]

Есть ли эффективный способ сделать это с помощью numpy? Все это критично для производительности.

Спасибо за вашу помощь!

Answer 1

Одним из способов будет использование np.minimum.accumulate вдоль этой оси -

np.minimum.accumulate(a<=threshold,axis=0)

Пробный прогон -

In [515]: a
Out[515]: array([1, 2, 2, 1, 2, 5, 2, 7, 1, 2])

In [516]: threshold = 3

In [518]: print np.minimum.accumulate(a<=threshold,axis=0)
[ True  True  True  True  True False False False False False]

Другой с порогом и затем нарезкой для1D массивы -

out = a<=threshold
if ~out.all():
    out[out.argmin():] = 0

Answer 2

Вот еще один подход, использующий 1st discrete difference:

In [126]: threshold = 3
In [127]: mask = np.diff(a, prepend=a[0]) < threshold

In [128]: mask[mask.argmin():] = False

In [129]: mask
Out[129]: 
array([ True,  True,  True,  True,  True, False, False, False, False,
       False])