Как использовать NumPy найти последовательные элементы и индекс при некоторых условиях?

Question

Предположим, у меня есть два numpy.array ():

a = np.array([1,2,3,4,5,6,7,8,9,10,11,12]) #index

b = np.array([1,1,4,3,3,3,5,2,2,2,2,6])    #element

Как мне эффективно найти все последующие элементы, число последовательных элементов> = 3, start_index и end_index? Я надеюсь, что результат:

c:  array([4,8])  #start_index
d:  array([6,11]) #end_index
e:  array([3,2])  #consequent same values whose number of consequent same element >=3

Answer 1

Найти границы группы:

sw, = np.where(np.diff(b, prepend=1, append=1))

Фильтр для условия:

cm, = np.where(np.diff(sw)>=3)

Получить соответствующие индексы

c = a[sw[cm]]
d = a[sw[cm+1]-1]

и значения

e = b[sw[cm]]

Полюбуйтесь

c
# array([4, 8])
d
# array([ 6, 11])
e
# array([3, 2])