条件下的 Python block 数组

Question

我想在特定条件下从数组中创建 block 。假设我想要所有不为空的 block 。

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]])

现在我知道第一个 block 从索引 0 开始并包含 [1]

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

它不必看起来像这样，但我不知道更好。到目前为止，我已经用 for 循环完成了。但每次我看我的代码时，它都让我很沮丧 :)。

有没有更好的方法(Numpy magic)？

Answer 1

试穿这件尺码:)

使用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])]