Python чанк массив при условии

Question

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

array = np.random.randint(2, size=10)
# array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0])

out = chunks(array)
# array([[0,[1]],[4,[1, 1]],[8,[1]])

Теперь я знаю, что первый чанк начинает индекс 0 и содержит [1]

0: 0 -> [1]
1: 4 -> [1,1]
2: 8 -> [1]

Это не обязательноПохоже на это, но я не знал ничего лучше.До сих пор я сделал это с помощью петель.Но каждый раз, когда я просматриваю свой код, меня это угнетает:).

Есть ли лучший способ (магия Numpy)?

Answer 1

Вот один из способов -

def group_elems(a, nullval=0):
    mask = a!=nullval
    mask_ext = np.r_[False, mask, False]
    idx = np.flatnonzero(mask_ext[1:] != mask_ext[:-1])
    cut_idx = (idx[1::2] - idx[::2]).cumsum()
    return dict(zip(idx[::2], np.split(a[mask], cut_idx[:-1])))

Образцы прогонов -

In [62]: a = np.array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0])

In [63]: group_elems(a, nullval=0)
Out[63]: {0: array([1]), 4: array([1, 1]), 8: array([1])}

In [64]: a = np.array([1, 0, 0, 0, 2, 3, 0, 0, 4, 0])

In [65]: group_elems(a, nullval=0)
Out[65]: {0: array([1]), 4: array([2, 3]), 8: array([4])}

Answer 2

Попробуйте этот размер:)

Использование np.where для фильтрации, np.diff и np.split для группировки последовательных индексов.

import numpy as np
array = np.array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0])
idx_matching_1 = np.where(array == 1)[0]
consecutive_groups_of_1_idx = np.split(idx_matching_1, np.where(np.diff(idx_matching_1) != 1)[0] + 1)
print(consecutive_groups_of_1_idx)


[array([0]), array([4, 5]), array([8])]