T.Poe
T.Poe

Reputation: 2089

Splitting a numpy array into two subsets of different sizes

I have a numpy array of a shape (400, 3, 3, 3) and I want to split it into two parts, so I would get arrays like (100, 3, 3, 3) and (300, 3, 3, 3).

I was playing with numpy split methods, e.g.:

subsets = np.array_split(arr, 2)

which gives me what I want, but it divides the original array into two halves the same size and I don't know how to specify these sizes. It'd be probably easy with some indexing (I guess) but I'm not sure how to do it.

Upvotes: 5

Views: 6154

Answers (1)

cs95
cs95

Reputation: 403218

As mentioned in my comment, you can use the Ellipsis notation to specify all axes:

x, y = arr[:100, ...], arr[100:, ...]

Upvotes: 6

Related Questions