[FIXED] extract blocks of columns (as seperated subarrays) indicated by 1D binary array

Issue

Based on a 1D binary mask, for example, np.array([0,0,0,1,1,1,0,0,1,1,0]), I would like to extract the columns of another array, indicated by the 1's in the binary mask, as as sub-arrays/separate blocks, like [9, 3.5, 7]) and [2.8, 9.1] (I am just making up the numbers to illustrate the point).

So far what I have (again just as a demo to illustrate what my goal is, not the data where this operation will be performed):

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

The output is:

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

What I would like to have in the end is:

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

Do you know how to implement it, preferably in pytorch, but numpy functions are also fine. A million thanks in advance!!

Solution

You can construct your tensor of slice indices with your approach. Only thing is you were missing the indices for the position of the end of each slice. You can do something like:

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

Then apply tensor_split and slice to only keep every other element:

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

Answered By - Ivan