基于标签对NumPy数组进行分区

Question

我有数据集A和B. A是一个形状为[169594,22]的矩阵 B是一个形状为[169594,1]的矩阵 B由（0,1,2,3,4,5）组成，它是A每行的标签。

所以，我想将A的数据分成每个标签。

所以我的代码如下。

我是Python的初学者，所以这段代码不起作用。

如果此代码运行良好，预期结果如下。

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

我该如何解决这个问题？ 谢谢！

Answer 1

您可以将B重塑为1D数组，然后在boolean indexing上使用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, :]

或者，更好的是，将隔离的项目保存在列表中。

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