Numpy Split array to sub arrays by specific number of elements-CodePudding

for example

I have a point list

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

and I have another array represents the number of elements

b = np.array([2,0,3,5])

how can I split array a according the number of elements of array b so that I can get the output

[[[0,0,0],[1,1,1]],
 [],
 [[2,2,2],[3,3,3],[4,4,4]],
 [[5,5,5],[6,6,6],[7,7,7],[8,8,8],[9,9,9]]]

CodePudding user response：

You can use numpy.split using cumsum on b to get the split points:

out = np.split(a, b.cumsum()[:-1])

output:

[array([[0, 0, 0],
        [1, 1, 1]]),
 array([], shape=(0, 3), dtype=int64),
 array([[2, 2, 2],
        [3, 3, 3],
        [4, 4, 4]]),
 array([[5, 5, 5],
        [6, 6, 6],
        [7, 7, 7],
        [8, 8, 8],
        [9, 9, 9]])]

If you want lists:

out = [x.tolist() for x in np.split(a, b.cumsum()[:-1])]

output:

[[[0, 0, 0], [1, 1, 1]],
 [],
 [[2, 2, 2], [3, 3, 3], [4, 4, 4]],
 [[5, 5, 5], [6, 6, 6], [7, 7, 7], [8, 8, 8], [9, 9, 9]]]

intermediate:

b.cumsum()[:-1]
# array([2, 2, 5])