使用另一个标签数组过滤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的图像索引。