How to split a numpy array with different array lengths (unequal array lengths)-CodePudding

Given an example array (or list), is there a way to split the array into different lengths? Here is desired input & output such that:

import numpy as np

# Input array
data = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17])

# Desired output splited arrays
[array([0, 1, 2, 3, 4, 5, 6, 7]), array([8, 9, 10, 11, 12, 13, 14, 15, 16, 17])]

I want to get the corresponding output, but it doesn't work, so I ask a question.

CodePudding user response：

Hard to know what you want exactly, but assuming you want 8 items for the first list, use numpy.array_split:

data = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17])

out = np.array_split(data, [8])

output:

[array([0, 1, 2, 3, 4, 5, 6, 7]),
 array([ 8,  9, 10, 11, 12, 13, 14, 15, 16, 17])]

To be more generic, you can use a list of sizes and process it with numpy.cumsum:

sizes = np.array([4,5,1,3])
out = np.array_split(data, np.cumsum(sizes))

output:

[array([0, 1, 2, 3]),           # 4 items
 array([4, 5, 6, 7, 8]),        # 5 items
 array([9]),                    # 1 item
 array([10, 11, 12]),           # 3 items
 array([13, 14, 15, 16, 17])]   # remaining items