Partitioning a NumPy array based on labels

Question

I have data sets A and B. A is a matrix which shape is [169594, 22] B is a matrix which shape is [169594,1] B consists of (0, 1, 2, 3, 4, 5) which is the label of each row of A.

So, I would like to separate Data of A into each labels.

So my code is as below.

I am beginner in Python, so this code is not work.

If this code is work well, the expected result is as below.

aa[xxx, 22]
bb[xxx, 22]
cc[xxx, 22]
dd[xxx, 22]
ee[xxx, 22]
ff[xxx, 22] 

How can I solve this problem? Thank you!

Answer 1

You could reshape B into a 1D array and then use boolean indexing on A .

B = B.reshape(-1, )
aa = A[B == 0, :]
bb = A[B == 1, :]
cc = A[B == 2, :]
dd = A[B == 3, :]
ee = A[B == 4, :]
ff = A[B == 5, :]

Or, even better, keep segregated items in a list.

l = []
for i in range(6):
    l.append(A[B == i, :]