提取以1D二进制数组表示的列块(作为分隔的子数组)

Question

基于一维二进制掩码，例如，np.array([0,0,0,1,1,1,0,0,1,1,0])，我想提取另一个数组的列，由二进制掩码中的1表示，作为子数组/单独的块，如[9, 3.5, 7]) and [2.8, 9.1](我只是编造数字来说明这一点)。

到目前为止，我所拥有的(再次作为演示来说明我的目标是什么，而不是执行该操作的数据):

arr = torch.from_numpy(np.array([0,0,0,1,1,1,0,0,1,1,0]))
split_idx = torch.where(torch.diff(arr) == 1)[0]+1
torch.tensor_split(arr, split_idx.tolist())

输出为:

(tensor([0, 0, 0]),
tensor([1, 1, 1]),
tensor([0, 0]),
tensor([1, 1]),
tensor([0]))

最后我想要的是:

(tensor([1, 1, 1]),
tensor([1, 1]))

你知道如何实现它吗，最好是在pytorch中，但numpy函数也很好。提前万分感谢!!

Answer 1

你可以用你的方法构造你的切片索引张量。唯一的问题是你漏掉了每个切片末端位置的索引。你可以这样做:

>>> slices = arr.diff().abs().nonzero().flatten()+1
tensor([ 3,  6,  8, 10])

然后应用tensor_split和slice只保留每一个其他元素:

>>> torch.tensor_split(arr, slices)[1::2]
(tensor([1, 1, 1]), tensor([1, 1]))