How to split a numpy array knowing the size of each subarray

Question

I want to split a numpy array, along the first axis, into subarrays of unequal size. I have checked numpy.split but it seems that I can only pass indices instead of size (number of rows for each subarray).

for example:

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

should yield:

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

Answer 1

Let the cutting intervals be -

cut_intvs = [2,1,2]

Then, use np.cumsum to detect positions of cuts -

cut_idx = np.cumsum(cut_intvs)

Finally, use those indices with np.split to cut the input array along the first axis and ignore the last cut to get the desired output, like so -

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

Sample run for explanation -

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