Numpy多维数组布尔索引问题：按一维数组条件筛选元素

这个问题我之前也碰到过，核心是Numpy布尔索引的维度匹配问题，咱们一步步来解决：

问题原因分析 #

你的y是(1,100)的二维数组，当你执行y==0时，得到的是一个同样(1,100)的布尔数组。直接用这个数组去索引X（形状(3,100)）时，Numpy会默认把布尔数组当作行索引来处理，但你的X只有3行，显然和(1,100)的掩码维度不匹配，所以会出错。

我们真正需要的是按列来筛选：因为X的每一列对应y中的一个元素，所以要在X的第二个维度（列维度）上应用掩码。

可行解决方案 #

这里有两种简单的方法，都能完美解决你的问题：

方法1：将y压缩为一维数组

用squeeze()或者flatten()把y从(1,100)转换成(100,)的一维数组，这样生成的掩码就是一维的，能直接匹配X的列维度：

import numpy as np

# 示例数据
y = np.random.randint(0, 2, (1, 100))
X = np.random.rand(3, 100)

# 筛选y=0对应的X列
mask_0 = y.squeeze() == 0
X_0 = X[:, mask_0]  # 结果形状为(3, n)，n是y中0的个数

# 筛选y=1对应的X列
mask_1 = y.squeeze() == 1
X_1 = X[:, mask_1]

方法2：直接取y的第一行

因为y只有一行，直接通过索引y[0]获取一维的标签数组，再生成掩码：

# 筛选y=0对应的X列
mask_0 = y[0] == 0
X_0 = X[:, mask_0]

# 筛选y=1对应的X列
mask_1 = y[0] == 1
X_1 = X[:, mask_1]

验证结果 #

你可以通过打印X_0.shape和X_1.shape来确认结果：比如如果y中有40个0和60个1，那么X_0的形状会是(3,40)，X_1是(3,60)，完全符合你的需求。

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