根據標簽對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整形為一維數組，然后在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, :]