如何使用np.where过滤条件含不同形状数组的numpy数组

核心误区说明 #

你当前对np.where()的适用场景存在混淆：

• 单参数调用np.where(cond)时，返回值是满足条件的元素索引组成的元组，这部分你已经正确获取到了，也就是代码里的Y_hex
• 三参数调用np.where(cond, x, y)是元素级值替换接口，返回数组形状和cond完全一致，仅会替换位置上的取值，不会删除元素、改变数组维度，完全不适合当前「筛选保留部分样本、压缩第一维长度」的需求，不需要硬凑这个接口的参数。

正确实现（一行完成过滤） #

你不需要写Python层的for循环，直接用NumPy原生的花式索引，配合已经拿到的索引数组，就可以一次性完成两个数组的过滤：

# 拿到符合Y<16的样本索引
keep_idx = np.where(Y < 16)[0]
# 同时过滤X和Y，一行完成
X_hex, Y_hex = X[keep_idx], Y[keep_idx]

如果要更简洁，还可以跳过np.where直接用布尔掩码索引，性能完全一致：

cond = Y < 16
X_hex, Y_hex = X[cond], Y[cond]

性能对比 #

• 上述两种NumPy原生索引实现，性能远高于Python for循环：索引逻辑在C层执行，没有Python循环的逐行解释开销，在你当前11.28万样本的数据集上，运行速度通常是手写for循环的数十倍到上百倍，内存利用效率也更高
• 强行使用三参数np.where()实现过滤反而会产生额外开销：因为三参数接口要求返回和原数组等长的结果，你还需要额外做切片删除不符合条件的元素，属于冗余操作，性能比直接索引差。

校验提示：过滤完成后可打印两个数组的shape，第一维长度会保持一致，且Y_hex的取值范围固定在0~15，正好覆盖0-9数字、A-F大写字母共16类十六进制字符，符合你的场景需求。

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