如何拆分一个numpy数组知道每个子数组的大小

Question

我想将一个 numpy 数组沿第一个轴拆分为大小不等的子数组。我检查了 numpy.split 但似乎我只能传递索引而不是大小（每个子数组的行数）。

例如：

arr = np.array([[1,2], [3,4], [5,6], [7,8], [9,10]])

应该让步：

arr.split([2,1,2]) = [array([[1,2], [3,4]]), array([5,6]), array([[7,8], [9,10]])]

Answer 1

让切割间隔为 -

cut_intvs = [2,1,2]

然后，使用np.cumsum 检测切割位置-

cut_idx = np.cumsum(cut_intvs)

最后，使用带有np.split 的索引沿第一个轴切割输入数组并忽略最后一个切割以获得所需的输出，就像这样 -

np.split(arr,np.cumsum(cut_intvs))[:-1]

示例运行解释 -

In [55]: arr                     # Input array
Out[55]: 
array([[ 1,  2],
       [ 3,  4],
       [ 5,  6],
       [ 7,  8],
       [ 9, 10]])

In [56]: cut_intvs = [2,1,2]    # Cutting intervals

In [57]: np.cumsum(cut_intvs)   # Indices at which cuts are to happen along axis=0 
Out[57]: array([2, 3, 5])

In [58]: np.split(arr,np.cumsum(cut_intvs))[:-1]  # Finally cut it, ignore last cut
Out[58]: 
[array([[1, 2],
        [3, 4]]), array([[5, 6]]), array([[ 7,  8],
        [ 9, 10]])]