使用另一個標簽數組過濾numpy數組

Question

給定兩個numpy數組，即：

images.shape: (60000, 784) # An array containing 60000 images
labels.shape: (60000, 10)  # An array of labels for each image

labels每一行在特定索引處都包含1 ，以指示images相關示例的類別。 （因此[0 0 1 0 0 0 0 0 0 0]表示該示例屬於Class 2（假設我們的類索引從0開始）。

我試圖有效地分離images以便可以一次操縱屬於特定類的所有圖像。 最明顯的解決方案是使用for循環（如下所示）。 但是，我不確定如何過濾images ，以便僅返回帶有適當labels images 。

for i in range(0, labels.shape[1]):
  class_images = # (?) Array containing all images that belong to class i

順便說一句，我也想知道是否還有更有效的方法可以消除對for循環的使用。

Answer 1

一種方法是將標簽數組轉換為bool並將其用於索引：

classes = []
blabels = labels.astype(bool)
for i in range(10):
    classes.append(images[blabels[:, i], :])

或者作為使用列表理解的單行代碼：

classes = [images[l.astype(bool), :] for l in labels.T]

Answer 2

_classes= [[] for x in range(10)]
for image_index , element in enumerate(labels):
    _classes[element.index(1)].append(image_index)

例如，_classes [0]將包含分類為class0的圖像索引。