如何用NumPy按y数组条件从x数组抽取样本构建MNIST数字0分类训练集

实现方案 #

使用NumPy原生布尔索引和随机抽样接口即可快速完成需求，无需手动遍历数组，完整实现代码如下：

import numpy as np

# 可根据需求调整每类抽取的样本数n
n = 1000

# 1. 筛选所有标签为0和非0的样本
x_zero = x_train[y_train == 0]
x_non_zero = x_train[y_train != 0]

# 2. 分别从两类中随机抽取n个不重复样本
# 若允许重复抽样可将replace改为True
selected_zero = x_zero[np.random.choice(x_zero.shape[0], size=n, replace=False)]
selected_non_zero = x_non_zero[np.random.choice(x_non_zero.shape[0], size=n, replace=False)]

# 3. 合并得到最终训练数组
custom_array = np.concatenate([selected_zero, selected_non_zero], axis=0)

# 可选：打乱数组顺序，避免前半段全为0类样本、后半段全为非0类样本
np.random.shuffle(custom_array)

补充说明 #

如果需要同步生成适配AdaBoost二分类任务的标签数组，可以额外添加以下代码：

# 0类标签设为1，非0类标签设为-1（符合AdaBoost二分类的标签规范）
custom_y = np.concatenate([np.ones(n), -np.ones(n)], axis=0)

# 同步打乱特征和标签，保证二者对应
shuffle_index = np.random.permutation(custom_array.shape[0])
custom_array = custom_array[shuffle_index]
custom_y = custom_y[shuffle_index]

MNIST训练集中0类样本共有约6000个，只要n不超过6000即可正常运行，无需担心样本量不足的问题。

内容的提问来源于stack exchange，提问作者MrFrenzoid