Получить группы последовательных элементов массива NumPy на основе условия

Question

У меня есть массив NumPy следующим образом:

import numpy as np
a = np.array([1, 4, 2, 6, 4, 4, 6, 2, 7, 6, 2, 8, 9, 3, 6, 3, 4, 4, 5, 8])

и постоянное число b = 6

Основано на предыдущем вопросе Я могу сосчитать число c, которое определяется числом раз, которое элементы в a меньше, чем b 2 или более раз подряд.

from itertools import groupby
b = 6
sum(len(list(g))>=2 for i, g in groupby(a < b) if i)

, поэтому в этом примере c == 3

Теперь я хотел бы выводить массив каждый раз, когда выполняется условие, а не подсчитывать, сколько раз условие выполняется.

Так что в этом примере правильным выводом будет:

array1 = [1, 4, 2]
array2 = [4, 4]
array3 = [3, 4, 4, 5]

с тех пор:

1, 4, 2, 6, 4, 4, 6, 2, 7, 6, 2, 8, 9, 3, 6, 3, 4, 4, 5, 8  # numbers in a
1, 1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0  # (a<b)
^^^^^^^-----^^^^-----------------------------^^^^^^^^^^---  # (a<b) 2+ times consecutively
   1         2                                    3

До сих пор я пробовал разные варианты:

np.isin((len(list(g))>=2 for i, g in groupby(a < b)if i), a)

и

np.extract((len(list(g))>=2 for i, g in groupby(a < b)if i), a)

Но ни один из них не достиг того, что я ищу.Может кто-нибудь указать мне правильные инструменты Python для вывода различных массивов, удовлетворяющих моему условию?

Answer 1

При измерении производительности мой другой ответ я заметил, что, хотя оно было быстрее, чем решение Остина (для массивов длины <15000), его сложность не была линейной.</p>

На основании этого ответа Я нашел следующее решение, используя np.split, что более эффективно, чем оба ранее добавленных ответа:

array = np.append(a, -np.inf)  # padding so we don't lose last element
mask = array >= 6  # values to be removed
split_indices = np.where(mask)[0]
for subarray in np.split(array, split_indices + 1):
    if len(subarray) > 2:
        print(subarray[:-1])

дает:

[1. 4. 2.]
[4. 4.]
[3. 4. 4. 5.]

Производительность *:

^{* Измерено perfplot}

Answer 2

Эта задача очень похожа на маркировку изображения , но, в вашем случае, она одномерна. Библиотека SciPy предоставляет некоторые полезные функции для обработки изображений, которые мы могли бы использовать здесь:

import numpy as np
from scipy.ndimage import (binary_dilation,
                           binary_erosion,
                           label)

a = np.array([1, 4, 2, 6, 4, 4, 6, 2, 7, 6, 2, 8, 9, 3, 6, 3, 4, 4, 5, 8])
b = 6  # your threshold
min_consequent_count = 2

mask = a < b
structure = [False] + [True] * min_consequent_count  # used for erosion and dilation
eroded = binary_erosion(mask, structure)
dilated = binary_dilation(eroded, structure)
labeled_array, labels_count = label(dilated)  # labels_count == c

for label_number in range(1, labels_count + 1):  # labeling starts from 1
    subarray = a[labeled_array == label_number]
    print(subarray)

дает:

[1 4 2]
[4 4]
[3 4 4 5]

---

Пояснение:

1. mask = a < b возвращает логический массив со значениями True, где элементы меньше порога b:

   array([ True,  True,  True, False,  True,  True, False,  True, False,
          False,  True, False, False,  True, False,  True,  True,  True,
           True, False])
   

2. Как видите, результат содержит несколько True элементов, у которых нет других True соседей вокруг них. Для их устранения мы могли бы использовать бинарную эрозию . Я использую scipy.ndimage.binary_erosion для этой цели. Его параметр по умолчанию structure не подходит для наших нужд, поскольку он также удалит два последовательных значения True, поэтому я создаю свое собственное:

   >>> structure = [False] + [True] * min_consequent_count
   >>> structure
   [False, True, True]
   >>> eroded = binary_erosion(mask, structure)
   >>> eroded
   array([ True,  True, False, False,  True, False, False, False, False,
          False, False, False, False, False, False,  True,  True,  True,
          False, False])
   

3. Нам удалось удалить отдельные значения True, но нам нужно получить начальную конфигурацию для других групп. Для этого мы используем двоичное расширение с тем же structure:

   >>> dilated = binary_dilation(eroded, structure)
   >>> dilated
   array([ True,  True,  True, False,  True,  True, False, False, False,
          False, False, False, False, False, False,  True,  True,  True,
           True, False])
   

   ^{Документы для binary_dilation: ссылка .}

4. И, наконец, мы помечаем каждую группу как scipy.ndimage.label:

   >>> labeled_array, labels_count = label(dilated)
   >>> labeled_array
   array([1, 1, 1, 0, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 3, 3, 3, 0])
   >>> labels_count
   3
   

   Вы можете видеть, что labels_count совпадает со значением c - числом групп в вопросе. Отсюда вы можете просто получить подгруппы путем логического индексирования:

   >>> a[labeled_array == 1]
   array([1, 4, 2])
   >>> a[labeled_array == 3]
   array([3, 4, 4, 5])
   

Answer 3

Используйте groupby и возьмите группы:

from itertools import groupby

lst = []
b = 6
for i, g in groupby(a, key=lambda x: x < b):
    grp = list(g)
    if i and len(grp) >= 2:
        lst.append(grp)

print(lst)

# [[1, 4, 2], [4, 4], [3, 4, 4, 5]]