Вот один метод, использующий .outer
и np.split
. Я сделал этот пример немного более интересным.
data = np.linspace(1, 5, 10).astype(int)
test_elements = np.array([1, 2, 4, 6])
y, x = np.where(np.equal.outer(test_elements,data))
np.split(x, y.searchsorted(np.arange(1,test_elements.size)))
# [array([0, 1, 2]), array([3, 4]), array([7, 8]), array([], dtype=int64)]
Некоторые объяснения:
np.equal.outer(test_elements,data)
совпадает с
test.elements[:,None] == data[None,:]
Итак, это двумерный логический массив, строки которого равны логическим массивам test_elem == data
, встречающимся в вашем цикле for.
where
возвращает два индексных массива из этого, по одному для каждой координаты. x
- это быстро меняющаяся координата, равная значениям, которые возвращает 1d where
вашего цикла for, но все они склеены в один длинный вектор. y
- это медленная координата изменения, ее значения упорядочены и могут использоваться для группировки / разделения x
. searchsorted
, возможно, не самый эффективный способ сделать это, но он прост и правильно обрабатывает пустые строки.