基於標籤對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, :]