如何在知道每个子数组的大小的情况下拆分 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]])]