Numpy 数组分割

Question

我有一个numpy 数组

import numpy as np

arr = np.arange(20).reshape(2,10)
arr[1,:] = 0
arr[1,2] = arr[1,5] = arr[1,7] = 1
print(arr)
>>>[[0 1 2 3 4 5 6 7 8 9]
>>> [0 0 1 0 0 1 0 1 0 0]]

我想提取重叠数组，从1 开始，到下一个1 之后结束。 预期输出：

[[0 1 2 3]
 [0 0 1 0]]

[[2 3 4 5 6]
 [1 0 0 1 0]]

[[5 6 7 8]
 [1 0 1 0]]

[[7 8 9]
 [1 0 0]]

目前，我有一个基于索引的 for 循环，在 numpy 上下文中感觉很尴尬，并且还必须将第一个和最后一个段视为特殊情况：

arr[1,0] = 1
ind = list(np.where(arr[1,:]))[0]
print(ind)

for i, j in enumerate(ind):
    if not i:
        continue
    curr = np.copy(arr[:, ind[i-1]:j+2])
    print(curr) 
        
#last segment
curr = np.copy(arr[:, j:])
print(curr)

这种方法给了我想要的输出，但我不敢相信没有一种更简单的方法可以实现这一点（尽管这里的风滚草反应可能表明这一点）。如果有更简单的熊猫解决方案，那也很好。理想情况下，输出是这些数组或类似数据结构的列表；输出数组不必单独返回。

Answer 1

有一部分解决方案，我最喜欢的，并不复杂：

split_idx = np.flatnonzero(arr[1]) + 2
>>> np.split(arr, split_idx, axis=1)
[array([[0, 1, 2, 3],
        [0, 0, 1, 0]]),
 array([[4, 5, 6],
        [0, 1, 0]]),
 array([[7, 8],
        [1, 0]]),
 array([[9],
        [0]])]

但是有两点表明针对此问题的任何numpyic 方法的设计不佳：

• 您不得不使用不是为numpy 设计的不同形状的列表。所以np.split 很慢。
• 您不能一次性循环一个数组。内部物品的开头需要额外插入。