基于标签对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阵列，然后在A上使用boolean indexing。

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, :]