Numpy: разделить массив на группы, которые удовлетворяют условию

Question

Я пытаюсь разбить массив с кусочками на куски, которые подчиняются определенному условию. Например, если a = np.array([ 1, 2, 3, 4, 10, 9, 8, 7, 7, 7, 8, 9, 10]) вернуть [array([3, 4]), array([7, 7, 7])], так как эти куски подчиняются условию np.logical_and(a>2,a<8).

Решение, которое я придумал, работает, но это неудобно, так как я применяю условие сначала для разделения массива, а затем снова для выбора нужных результатов.

import numpy as np
a = np.array([ 1,  2,  3,  4, 10,  9,  8,  7,  7,  7,  8,  9, 10])
cond = np.logical_and(a>2,a<8)
splits = np.split(a,np.where(np.diff(cond))[0]+1)
# this yields [array([1, 2]), array([3, 4]), array([10,  9,  8]), array([7, 7, 7]), array([ 8,  9, 10])]
result = [x for x in splits if np.logical_and(x>2,x<8).any()]
# [array([3, 4]), array([7, 7, 7])]

Есть ли способ применить одно и то же условие к обоим, чтобы, если бы я захотел изменить условие, я мог изменить его один раз в своем коде, а не дважды? Или я слишком усложняю вещи, когда есть гораздо более простой, более питонический способ сделать это?

Answer 1

Вот одна опция, подмножество массива перед разбиением, чтобы исключить цикл for :

cond = (a > 2) & (a < 8)
np.split(a[cond], np.where(np.diff(np.where(cond)[0]) > 1)[0] + 1)
# [array([3, 4]), array([7, 7, 7])]