Home > Software design >  Slice numpy array into groups
Slice numpy array into groups

Time:11-18

We have the following numpy array:

A = [[0,1,2,3],
[2,3,4,5],
[1,2,3,4],
[1,1,1,1],
[2,2,2,2],
[3,3,3,3],
[0,1,2,3],
[2,3,4,5],
[1,2,3,4],
 [10,10,10,10],
[20,20,20,20],
[30,30,30,30]]

I would like to create two new arrays:

B = [[0,1,2,3],
[2,3,4,5],
[1,2,3,4],
[0,1,2,3],
[2,3,4,5],
[1,2,3,4]]

and

C = [[1,1,1,1],
[2,2,2,2],
[3,3,3,3],
 [10,10,10,10],
[20,20,20,20],
[30,30,30,30]]

Basically I would like to split array A into two new arrays where B takes in groups of 3 rows, and array C takes the next groups of 3 rows.

CodePudding user response:

There are probabily many different approaches to solve this, here is one:

import numpy as np
A = np.array([[0,1,2,3],
[2,3,4,5],
[1,2,3,4],
[1,1,1,1],
[2,2,2,2],
[3,3,3,3],
[0,1,2,3],
[2,3,4,5],
[1,2,3,4],
 [10,10,10,10],
[20,20,20,20],
[30,30,30,30]])

indices = np.reshape(np.arange(A.shape[0]),(-1,3))

B = A[indices[::2].flatten()]
C = A[indices[1::2].flatten()]

CodePudding user response:

Let's try:

tmp = A.reshape(-1, 3, A.shape[1])
B = tmp[::2]
C = tmp[1::2]
  • Related