如何在知道每个子阵列大小的情况下拆分numpy数组

时间:2016-03-03 09:55:37

标签: python numpy

我想沿第一轴将一个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]])]

1 个答案:

答案 0 :(得分:4)

切割间隔为 -

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